<s> from helpers import unittest <EOL> import mock <EOL> import sys <EOL> from helpers import with_config <EOL> from luigi import notifications <EOL> from luigi import configuration <EOL> from luigi . scheduler import CentralPlannerScheduler <EOL> from luigi . worker import Worker <EOL> from luigi import six <EOL> import luigi <EOL> class TestEmail ( unittest . TestCase ) : <EOL> def testEmailNoPrefix ( self ) : <EOL> self . assertEqual ( "<STR_LIT>" , notifications . _prefix ( '<STR_LIT>' ) ) <EOL> @ with_config ( { "<STR_LIT>" : { "<STR_LIT>" : "<STR_LIT>" } } ) <EOL> def testEmailPrefix ( self ) : <EOL> self . assertEqual ( "<STR_LIT>" , notifications . _prefix ( '<STR_LIT>' ) ) <EOL> class TestException ( Exception ) : <EOL> pass <EOL> class TestTask ( luigi . Task ) : <EOL> foo = luigi . Parameter ( ) <EOL> bar = luigi . Parameter ( ) <EOL> class FailSchedulingTask ( TestTask ) : <EOL> def requires ( self ) : <EOL> raise TestException ( '<STR_LIT>' ) <EOL> def run ( self ) : <EOL> pass <EOL> def complete ( self ) : <EOL> return False <EOL> class FailRunTask ( TestTask ) : <EOL> def run ( self ) : <EOL> raise TestException ( '<STR_LIT>' ) <EOL> def complete ( self ) : <EOL> return False <EOL> class ExceptionFormatTest ( unittest . TestCase ) : <EOL> def setUp ( self ) : <EOL> self . sch = CentralPlannerScheduler ( ) <EOL> def test_fail_run ( self ) : <EOL> task = FailRunTask ( foo = '<STR_LIT:foo>' , bar = '<STR_LIT:bar>' ) <EOL> self . _run_task ( task ) <EOL> def test_fail_run_html ( self ) : <EOL> task = FailRunTask ( foo = '<STR_LIT:foo>' , bar = '<STR_LIT:bar>' ) <EOL> self . _run_task_html ( task ) <EOL> def test_fail_schedule ( self ) : <EOL> task = FailSchedulingTask ( foo = '<STR_LIT:foo>' , bar = '<STR_LIT:bar>' ) <EOL> self . _run_task ( task ) <EOL> def test_fail_schedule_html ( self ) : <EOL> task = FailSchedulingTask ( foo = '<STR_LIT:foo>' , bar = '<STR_LIT:bar>' ) <EOL> self . _run_task_html ( task ) <EOL> @ with_config ( { '<STR_LIT>' : { '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' } } ) <EOL> @ mock . patch ( '<STR_LIT>' ) <EOL> def _run_task ( self , task , mock_send ) : <EOL> with Worker ( scheduler = self . sch ) as w : <EOL> w . add ( task ) <EOL> w . run ( ) <EOL> self . assertEqual ( mock_send . call_count , <NUM_LIT:1> ) <EOL> args , kwargs = mock_send . call_args <EOL> self . _check_subject ( args [ <NUM_LIT:0> ] , task ) <EOL> self . _check_body ( args [ <NUM_LIT:1> ] , task , html = False ) <EOL> @ with_config ( { '<STR_LIT>' : { '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT:html>' } } ) <EOL> @ mock . patch ( '<STR_LIT>' ) <EOL> def _run_task_html ( self , task , mock_send ) : <EOL> with Worker ( scheduler = self . sch ) as w : <EOL> w . add ( task ) <EOL> w . run ( ) <EOL> self . assertEqual ( mock_send . call_count , <NUM_LIT:1> ) <EOL> args , kwargs = mock_send . call_args <EOL> self . _check_subject ( args [ <NUM_LIT:0> ] , task ) <EOL> self . _check_body ( args [ <NUM_LIT:1> ] , task , html = True ) <EOL> def _check_subject ( self , subject , task ) : <EOL> self . assertIn ( str ( task ) , subject ) <EOL> def _check_body ( self , body , task , html = False ) : <EOL> if html : <EOL> self . assertIn ( '<STR_LIT>' . format ( task . task_family ) , body ) <EOL> self . assertIn ( '<STR_LIT>' , body ) <EOL> self . assertIn ( '<STR_LIT>' , body ) <EOL> for param , value in task . param_kwargs . items ( ) : <EOL> self . assertIn ( '<STR_LIT>' . format ( param , value ) , body ) <EOL> else : <EOL> self . assertIn ( '<STR_LIT>' . format ( task . task_family ) , body ) <EOL> self . assertIn ( '<STR_LIT>' , body ) <EOL> self . assertIn ( '<STR_LIT>' , body ) <EOL> for param , value in task . param_kwargs . items ( ) : <EOL> self . assertIn ( '<STR_LIT>' . format ( param , value ) , body ) <EOL> @ with_config ( { "<STR_LIT>" : { "<STR_LIT>" : "<STR_LIT>" } } ) <EOL> def testEmailRecipients ( self ) : <EOL> six . assertCountEqual ( self , notifications . _email_recipients ( ) , [ "<STR_LIT>" ] ) <EOL> six . assertCountEqual ( self , notifications . _email_recipients ( "<STR_LIT>" ) , [ "<STR_LIT>" , "<STR_LIT>" ] ) <EOL> six . assertCountEqual ( self , notifications . _email_recipients ( [ "<STR_LIT>" , "<STR_LIT>" ] ) , <EOL> [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] ) <EOL> @ with_config ( { "<STR_LIT>" : { } } , replace_sections = True ) <EOL> def testEmailRecipientsNoConfig ( self ) : <EOL> six . assertCountEqual ( self , notifications . _email_recipients ( ) , [ ] ) <EOL> six . assertCountEqual ( self , notifications . _email_recipients ( "<STR_LIT>" ) , [ "<STR_LIT>" ] ) <EOL> six . assertCountEqual ( self , notifications . _email_recipients ( [ "<STR_LIT>" , "<STR_LIT>" ] ) , <EOL> [ "<STR_LIT>" , "<STR_LIT>" ] ) <EOL> class NotificationFixture ( object ) : <EOL> """<STR_LIT>""" <EOL> sender = '<STR_LIT>' <EOL> subject = '<STR_LIT>' <EOL> message = """<STR_LIT>""" <EOL> recipients = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> image_png = None <EOL> notification_args = [ sender , subject , message , recipients , image_png ] <EOL> mocked_email_msg = '''<STR_LIT>''' <EOL> class TestSMTPEmail ( unittest . TestCase , NotificationFixture ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> sys . modules [ '<STR_LIT>' ] = mock . MagicMock ( ) <EOL> import smtplib <EOL> def tearDown ( self ) : <EOL> del sys . modules [ '<STR_LIT>' ] <EOL> @ with_config ( { "<STR_LIT>" : { "<STR_LIT>" : "<STR_LIT:False>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT:False>" } } ) <EOL> def test_sends_smtp_email ( self ) : <EOL> """<STR_LIT>""" <EOL> smtp_kws = { "<STR_LIT:host>" : "<STR_LIT>" , <EOL> "<STR_LIT:port>" : <NUM_LIT> , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : <NUM_LIT> } <EOL> with mock . patch ( '<STR_LIT>' ) as SMTP : <EOL> with mock . patch ( '<STR_LIT>' ) as generate_email : <EOL> generate_email . return_value . as_string . return_value = self . mocked_email_msg <EOL> notifications . send_email_smtp ( configuration . get_config ( ) , <EOL> * self . notification_args ) <EOL> SMTP . assert_called_once_with ( ** smtp_kws ) <EOL> SMTP . return_value . login . assert_called_once_with ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> SMTP . return_value . starttls . assert_called_once_with ( ) <EOL> SMTP . return_value . sendmail . assert_called_once_with ( self . sender , self . recipients , <EOL> self . mocked_email_msg ) <EOL> @ with_config ( { "<STR_LIT>" : { "<STR_LIT>" : "<STR_LIT:False>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT:True>" } } ) <EOL> def test_sends_smtp_email_without_tls ( self ) : <EOL> """<STR_LIT>""" <EOL> smtp_kws = { "<STR_LIT:host>" : "<STR_LIT>" , <EOL> "<STR_LIT:port>" : <NUM_LIT> , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : <NUM_LIT> } <EOL> with mock . patch ( '<STR_LIT>' ) as SMTP : <EOL> with mock . patch ( '<STR_LIT>' ) as generate_email : <EOL> generate_email . return_value . as_string . return_value = self . mocked_email_msg <EOL> notifications . send_email_smtp ( configuration . get_config ( ) , <EOL> * self . notification_args ) <EOL> SMTP . assert_called_once_with ( ** smtp_kws ) <EOL> self . assertEqual ( SMTP . return_value . starttls . called , False ) <EOL> SMTP . return_value . login . assert_called_once_with ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> SMTP . return_value . sendmail . assert_called_once_with ( self . sender , self . recipients , <EOL> self . mocked_email_msg ) <EOL> class TestSendgridEmail ( unittest . TestCase , NotificationFixture ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> sys . modules [ '<STR_LIT>' ] = mock . MagicMock ( ) <EOL> import sendgrid <EOL> def tearDown ( self ) : <EOL> del sys . modules [ '<STR_LIT>' ] <EOL> @ with_config ( { "<STR_LIT:email>" : { "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" } } ) <EOL> def test_sends_sendgrid_email ( self ) : <EOL> """<STR_LIT>""" <EOL> with mock . patch ( '<STR_LIT>' ) as SendgridClient : <EOL> notifications . send_email_sendgrid ( configuration . get_config ( ) , <EOL> * self . notification_args ) <EOL> SendgridClient . assert_called_once_with ( "<STR_LIT>" , "<STR_LIT>" , raise_errors = True ) <EOL> self . assertTrue ( SendgridClient . return_value . send . called ) <EOL> class TestSESEmail ( unittest . TestCase , NotificationFixture ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> sys . modules [ '<STR_LIT>' ] = mock . MagicMock ( ) <EOL> import boto3 <EOL> def tearDown ( self ) : <EOL> del sys . modules [ '<STR_LIT>' ] <EOL> @ with_config ( { } ) <EOL> def test_sends_ses_email ( self ) : <EOL> """<STR_LIT>""" <EOL> with mock . patch ( '<STR_LIT>' ) as boto_client : <EOL> with mock . patch ( '<STR_LIT>' ) as generate_email : <EOL> generate_email . return_value . as_string . return_value = self . mocked_email_msg <EOL> notifications . send_email_ses ( configuration . get_config ( ) , <EOL> * self . notification_args ) <EOL> SES = boto_client . return_value <EOL> SES . send_raw_email . assert_called_once_with ( <EOL> Source = self . sender , <EOL> Destinations = self . recipients , <EOL> RawMessage = { '<STR_LIT>' : self . mocked_email_msg } ) <EOL> class TestSNSNotification ( unittest . TestCase , NotificationFixture ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> sys . modules [ '<STR_LIT>' ] = mock . MagicMock ( ) <EOL> import boto3 <EOL> def tearDown ( self ) : <EOL> del sys . modules [ '<STR_LIT>' ] <EOL> @ with_config ( { } ) <EOL> def test_sends_sns_email ( self ) : <EOL> """<STR_LIT>""" <EOL> with mock . patch ( '<STR_LIT>' ) as res : <EOL> notifications . send_email_sns ( configuration . get_config ( ) , <EOL> * self . notification_args ) <EOL> SNS = res . return_value <EOL> SNS . Topic . assert_called_once_with ( self . recipients [ <NUM_LIT:0> ] ) <EOL> SNS . Topic . return_value . publish . assert_called_once_with ( <EOL> Subject = self . subject , Message = self . message ) <EOL> @ with_config ( { } ) <EOL> def test_sns_subject_is_shortened ( self ) : <EOL> """<STR_LIT>""" <EOL> long_subject = '<STR_LIT>' '<STR_LIT>' <EOL> with mock . patch ( '<STR_LIT>' ) as res : <EOL> notifications . send_email_sns ( configuration . get_config ( ) , <EOL> self . sender , long_subject , self . message , <EOL> self . recipients , self . image_png ) <EOL> SNS = res . return_value <EOL> SNS . Topic . assert_called_once_with ( self . recipients [ <NUM_LIT:0> ] ) <EOL> called_subj = SNS . Topic . return_value . publish . call_args [ <NUM_LIT:1> ] [ '<STR_LIT>' ] <EOL> self . assertTrue ( len ( called_subj ) <= <NUM_LIT:100> , <EOL> "<STR_LIT>" . format ( len ( called_subj ) ) ) <EOL> class Test_Notification_Dispatcher ( unittest . TestCase , NotificationFixture ) : <EOL> """<STR_LIT>""" <EOL> def check_dispatcher ( self , target ) : <EOL> """<STR_LIT>""" <EOL> expected_args = self . notification_args <EOL> with mock . patch ( '<STR_LIT>' . format ( target ) ) as sender : <EOL> notifications . send_email ( self . subject , self . message , self . sender , <EOL> self . recipients , image_png = self . image_png ) <EOL> self . assertTrue ( sender . called ) <EOL> call_args = sender . call_args [ <NUM_LIT:0> ] [ <NUM_LIT:1> : ] <EOL> self . assertEqual ( tuple ( expected_args ) , call_args ) <EOL> @ with_config ( { '<STR_LIT:email>' : { '<STR_LIT>' : '<STR_LIT:True>' , <EOL> '<STR_LIT:type>' : '<STR_LIT>' } } ) <EOL> def test_smtp ( self ) : <EOL> return self . check_dispatcher ( '<STR_LIT>' ) <EOL> @ with_config ( { '<STR_LIT:email>' : { '<STR_LIT>' : '<STR_LIT:True>' , <EOL> '<STR_LIT:type>' : '<STR_LIT>' } } ) <EOL> def test_ses ( self ) : <EOL> return self . check_dispatcher ( '<STR_LIT>' ) <EOL> @ with_config ( { '<STR_LIT:email>' : { '<STR_LIT>' : '<STR_LIT:True>' , <EOL> '<STR_LIT:type>' : '<STR_LIT>' } } ) <EOL> def test_sendgrid ( self ) : <EOL> return self . check_dispatcher ( '<STR_LIT>' ) <EOL> @ with_config ( { '<STR_LIT:email>' : { '<STR_LIT>' : '<STR_LIT:True>' , <EOL> '<STR_LIT:type>' : '<STR_LIT>' } } ) <EOL> def test_sns ( self ) : <EOL> return self . check_dispatcher ( '<STR_LIT>' ) </s>
