def test_poke_with_attachment_not_found(self, mock_imap_hook): mock_imap_hook.return_value.__enter__ = Mock( return_value=mock_imap_hook) mock_imap_hook.has_mail_attachment.return_value = False imap_attachment_sensor = ImapAttachmentSensor( conn_id='imap_test', attachment_name='test_attachment', task_id='check_for_attachment_on_mail_server_test', dag=None) self.assertFalse(imap_attachment_sensor.poke(context={})) mock_imap_hook.has_mail_attachment.assert_called_once_with( name='test_attachment', mail_folder='INBOX', check_regex=False)
def test_poke(self, has_attachment_return_value, mock_imap_hook): mock_imap_hook.return_value.__enter__ = Mock( return_value=mock_imap_hook) mock_imap_hook.has_mail_attachment.return_value = has_attachment_return_value has_attachment = ImapAttachmentSensor(**self.kwargs).poke(context={}) self.assertEqual(has_attachment, mock_imap_hook.has_mail_attachment.return_value) mock_imap_hook.has_mail_attachment.assert_called_once_with( name=self.kwargs['attachment_name'], check_regex=self.kwargs['check_regex'], mail_folder=self.kwargs['mail_folder'], mail_filter=self.kwargs['mail_filter'])