def test_synthesize_text_green_path(self, mock_text_to_speech_hook, mock_gcp_hook): mocked_response = Mock() type(mocked_response).audio_content = PropertyMock(return_value=b"audio") mock_text_to_speech_hook.return_value.synthesize_speech.return_value = mocked_response mock_gcp_hook.return_value.upload.return_value = True GcpTextToSpeechSynthesizeOperator( project_id=PROJECT_ID, gcp_conn_id=GCP_CONN_ID, input_data=INPUT, voice=VOICE, audio_config=AUDIO_CONFIG, target_bucket_name=TARGET_BUCKET_NAME, target_filename=TARGET_FILENAME, task_id="id", ).execute(context={"task_instance": Mock()}) mock_text_to_speech_hook.assert_called_once_with(gcp_conn_id="gcp-conn-id") mock_gcp_hook.assert_called_once_with(google_cloud_storage_conn_id="gcp-conn-id") mock_text_to_speech_hook.return_value.synthesize_speech.assert_called_once_with( input_data=INPUT, voice=VOICE, audio_config=AUDIO_CONFIG, retry=None, timeout=None ) mock_gcp_hook.return_value.upload.assert_called_once_with( bucket_name=TARGET_BUCKET_NAME, object_name=TARGET_FILENAME, filename=ANY )
def test_fakesock_socket_sendall_with_body_data_with_chunked_entry(POTENTIAL_HTTP_PORTS, old_socket, match_uriinfo): ("fakesock.socket#sendall should call real_sendall when not ") # Background: # Using a subclass of socket that mocks out real_sendall class MySocket(fakesock.socket): def real_sendall(self, data): raise AssertionError('should have never been called') matcher = Mock(name='matcher') info = Mock(name='info') httpretty.match_uriinfo.return_value = (matcher, info) # Using a mocked entry entry = Mock() entry.method = 'GET' entry.info.path = '/foo' entry.request.headers = { 'transfer-encoding': 'chunked', } entry.request.body = b'' # Given an instance of that socket socket = MySocket() socket._entry = entry # And that is is considered http socket.connect(('foo.com', 80)) # When I try to send data socket.sendall(b"BLABLABLABLA") # Then the entry should have that body httpretty.last_request.body.should.equal(b'BLABLABLABLA')
def test_missing_arguments( self, missing_arg, input_data, voice, audio_config, target_bucket_name, target_filename, mock_text_to_speech_hook, mock_gcp_hook, ): with self.assertRaises(AirflowException) as e: GcpTextToSpeechSynthesizeOperator( project_id="project-id", input_data=input_data, voice=voice, audio_config=audio_config, target_bucket_name=target_bucket_name, target_filename=target_filename, task_id="id", ).execute(context={"task_instance": Mock()}) err = e.exception self.assertIn(missing_arg, str(err)) mock_text_to_speech_hook.assert_not_called() mock_gcp_hook.assert_not_called()
class TestStatsWithAllowList(unittest.TestCase): def setUp(self): from airflow.settings import SafeStatsdLogger, AllowListValidator self.statsd_client = Mock() self.stats = SafeStatsdLogger( self.statsd_client, AllowListValidator("stats_one, stats_two")) def test_increment_counter_with_allowed_key(self): self.stats.incr('stats_one') self.statsd_client.incr.assert_called_once_with('stats_one', 1, 1) def test_increment_counter_with_allowed_prefix(self): self.stats.incr('stats_two.bla') self.statsd_client.incr.assert_called_once_with('stats_two.bla', 1, 1) def test_not_increment_counter_if_not_allowed(self): self.stats.incr('stats_three') self.statsd_client.assert_not_called()
def test_initialize_order(self, prepare_syspath, import_local_settings): """ Tests that import_local_settings is called after prepare_classpath """ mock = Mock() mock.attach_mock(prepare_syspath, "prepare_syspath") mock.attach_mock(import_local_settings, "import_local_settings") import airflow.settings airflow.settings.initialize() mock.assert_has_calls([call.prepare_syspath(), call.import_local_settings()])
def setUp(self): self.sentry = ConfiguredSentry() # Mock the Dag self.dag = Mock(dag_id=DAG_ID, params=[]) self.dag.task_ids = [TASK_ID] # Mock the task self.task = Mock(dag=self.dag, dag_id=DAG_ID, task_id=TASK_ID, params=[], pool_slots=1) self.task.__class__.__name__ = OPERATOR self.ti = TaskInstance(self.task, execution_date=EXECUTION_DATE) self.ti.operator = OPERATOR self.ti.state = STATE self.dag.get_task_instances = MagicMock(return_value=[self.ti]) self.session = Session()
def test_recognize_speech_green_path(self, mock_hook): mock_hook.return_value.recognize_speech.return_value = True GcpSpeechToTextRecognizeSpeechOperator( # pylint: disable=no-value-for-parameter project_id=PROJECT_ID, gcp_conn_id=GCP_CONN_ID, config=CONFIG, audio=AUDIO, task_id="id" ).execute(context={"task_instance": Mock()}) mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID) mock_hook.return_value.recognize_speech.assert_called_once_with( config=CONFIG, audio=AUDIO, retry=None, timeout=None )
def test_missing_audio(self, mock_hook): mock_hook.return_value.recognize_speech.return_value = True with self.assertRaises(AirflowException) as e: GcpSpeechToTextRecognizeSpeechOperator( # pylint: disable=no-value-for-parameter project_id=PROJECT_ID, gcp_conn_id=GCP_CONN_ID, config=CONFIG, task_id="id" ).execute(context={"task_instance": Mock()}) err = e.exception self.assertIn("audio", str(err)) mock_hook.assert_not_called()
def test_fakesock_socket_sendall_with_valid_requestline_2(POTENTIAL_HTTP_PORTS, httpretty, old_socket): ("fakesock.socket#sendall should create an entry if it's given a valid request line") matcher = Mock(name='matcher') info = Mock(name='info') httpretty.match_uriinfo.return_value = (matcher, info) httpretty.register_uri(httpretty.GET, 'http://foo.com/foobar') # Background: # using a subclass of socket that mocks out real_sendall class MySocket(fakesock.socket): def real_sendall(self, data, *args, **kw): raise AssertionError('should never call this...') # Given an instance of that socket socket = MySocket() # And that is is considered http socket.connect(('foo.com', 80)) # When I try to send data socket.sendall(b"GET /foobar HTTP/1.1\r\nContent-Type: application/json\r\n\r\n")
def test_py26_callback_response(): ("HTTPretty should call a callback function *once* and set its return value" " as the body of the response requests") def _request_callback(request, uri, headers): return [ 200, headers, "The {} response from {}".format(decode_utf8(request.method), uri) ] request_callback = Mock() request_callback.side_effect = _request_callback HTTPretty.register_uri(HTTPretty.POST, "https://api.yahoo.com/test_post", body=request_callback) requests.post("https://api.yahoo.com/test_post", {"username": "******"}) os.environ['STOP'] = 'true' expect(request_callback.call_count).equal(1)
def test_trigger_controller_dag(self): dag = self.dagbag.get_dag('example_trigger_controller_dag') target_dag = self.dagbag.get_dag('example_trigger_target_dag') target_dag.sync_to_db() dag_file_processor = DagFileProcessor(dag_ids=[], log=Mock()) task_instances_list = Mock() dag_file_processor._process_task_instances( target_dag, task_instances_list=task_instances_list) self.assertFalse(task_instances_list.append.called) job = BackfillJob(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_first_depends_on_past=True) job.run() dag_file_processor._process_task_instances( target_dag, task_instances_list=task_instances_list) self.assertTrue(task_instances_list.append.called)
def test_fakesock_socket_ssl(): ("fakesock.socket#ssl should take a socket instance and return itself") # Given a fake socket instance socket = fakesock.socket() # And a stubbed socket sentinel sentinel = Mock() # When I call `ssl` on that mock result = socket.ssl(sentinel) # Then it should have returned its first argument result.should.equal(sentinel)
def test_missing_config(self, mock_hook): mock_hook.return_value.recognize_speech.return_value = True with self.assertRaises(AirflowException) as e: GcpSpeechToTextRecognizeSpeechOperator( project_id=PROJECT_ID, gcp_conn_id=GCP_CONN_ID, audio=AUDIO, task_id="id").execute(context={"task_instance": Mock()}) err = e.exception self.assertIn("config", str(err)) mock_hook.assert_not_called()
def test_heartbeat_failed(self, mock_create_session): when = timezone.utcnow() - datetime.timedelta(seconds=60) with create_session() as session: mock_session = Mock(spec_set=session, name="MockSession") mock_create_session.return_value.__enter__.return_value = mock_session job = self.TestJob(None, heartrate=10, state=State.RUNNING) job.latest_heartbeat = when mock_session.commit.side_effect = OperationalError("Force fail", {}, None) job.heartbeat() self.assertEqual(job.latest_heartbeat, when, "attribute not updated when heartbeat fails")
def test_fake_ssl_socket_proxies_its_ow_socket(): ("FakeSSLSocket is a simpel wrapper around its own socket, " "which was designed to be a HTTPretty fake socket") # Given a sentinel mock object socket = Mock() # And a FakeSSLSocket wrapping it ssl = FakeSSLSocket(socket) # When I make a method call ssl.send("FOO") # Then it should bypass any method calls to its own socket socket.send.assert_called_once_with("FOO")
def test_execute(self, mock_tmp_file, mock_hive_hook, mock_samba_hook): type(mock_tmp_file).name = PropertyMock(return_value='tmp_file') mock_tmp_file.return_value.__enter__ = Mock(return_value=mock_tmp_file) context = {} Hive2SambaOperator(**self.kwargs).execute(context) mock_hive_hook.assert_called_once_with( hiveserver2_conn_id=self.kwargs['hiveserver2_conn_id']) mock_hive_hook.return_value.to_csv.assert_called_once_with( hql=self.kwargs['hql'], csv_filepath=mock_tmp_file.name, hive_conf=context_to_airflow_vars(context)) mock_samba_hook.assert_called_once_with( samba_conn_id=self.kwargs['samba_conn_id']) mock_samba_hook.return_value.push_from_local.assert_called_once_with( self.kwargs['destination_filepath'], mock_tmp_file.name)
def test_none_failed_or_skipped_tr_skipped(self): """ All success including all upstream skips trigger rule success """ ti = self._get_task_instance( TriggerRule.NONE_FAILED_OR_SKIPPED, upstream_task_ids=["FakeTaskID", "OtherFakeTaskID"]) dep_statuses = tuple(TriggerRuleDep()._evaluate_trigger_rule( ti=ti, successes=0, skipped=2, failed=0, upstream_failed=0, done=2, flag_upstream_failed=True, session=Mock())) self.assertEqual(len(dep_statuses), 0) self.assertEqual(ti.state, State.SKIPPED)
def test_request_string_representation(): ("HTTPrettyRequest should have a forward_and_trace-friendly " "string representation") # Given a request string containing a unicode encoded querystring headers = "\r\n".join([ 'POST /create HTTP/1.1', 'Content-Type: JPEG-baby', 'Host: blog.falcao.it' ]) # And a valid urlencoded body body = "foobar:\nlalala" # When I create a HTTPrettyRequest with that data request = HTTPrettyRequest(headers, body, sock=Mock(is_https=True)) # Then its string representation should show the headers and the body str(request).should.equal('<HTTPrettyRequest("POST", "https://blog.falcao.it/create", headers={\'Content-Type\': \'JPEG-baby\', \'Host\': \'blog.falcao.it\'}, body=14)>')
def test_all_success_tr_skip_flag_upstream(self): """ All-success trigger rule fails when some upstream tasks are skipped. The state of the ti should be set to SKIPPED when flag_upstream_failed is True. """ ti = self._get_task_instance( TriggerRule.ALL_SUCCESS, upstream_task_ids=["FakeTaskID", "OtherFakeTaskID"]) dep_statuses = tuple(TriggerRuleDep()._evaluate_trigger_rule( ti=ti, successes=1, skipped=1, failed=0, upstream_failed=0, done=2, flag_upstream_failed=True, session=Mock())) self.assertEqual(len(dep_statuses), 1) self.assertFalse(dep_statuses[0].passed) self.assertEqual(ti.state, State.SKIPPED)
def test_fakesock_socket_makefile(old_socket): ("fakesock.socket#makefile should set the mode, " "bufsize and return its mocked file descriptor") # Given a fake socket that has a mocked Entry associated with it socket = fakesock.socket() socket._entry = Mock() # When I call makefile() fd = socket.makefile(mode='rw', bufsize=512) # Then it should have returned the socket's own filedescriptor expect(fd).to.equal(socket.fd) # And the mode should have been set in the socket instance socket._mode.should.equal('rw') # And the bufsize should have been set in the socket instance socket._bufsize.should.equal(512) # And the entry should have been filled with that filedescriptor socket._entry.fill_filekind.assert_called_once_with(fd)
def setUp(self): from airflow.settings import SafeStatsdLogger, AllowListValidator self.statsd_client = Mock() self.stats = SafeStatsdLogger( self.statsd_client, AllowListValidator("stats_one, stats_two"))
def test_pod_mutation_v1_pod(self): with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK_V1_POD, "airflow_local_settings"): from airflow import settings settings.import_local_settings() # pylint: ignore from airflow.kubernetes.pod_launcher import PodLauncher self.mock_kube_client = Mock() self.pod_launcher = PodLauncher(kube_client=self.mock_kube_client) pod = pod_generator.PodGenerator(image="myimage", cmds=["foo"], namespace="baz", volume_mounts=[{ "name": "foo", "mountPath": "/mnt", "subPath": "/", "readOnly": True }], volumes=[{ "name": "foo" }]).gen_pod() sanitized_pod_pre_mutation = api_client.sanitize_for_serialization( pod) self.assertEqual( sanitized_pod_pre_mutation, { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'namespace': 'baz' }, 'spec': { 'containers': [{ 'args': [], 'command': ['foo'], 'env': [], 'envFrom': [], 'image': 'myimage', 'name': 'base', 'ports': [], 'volumeMounts': [{ 'mountPath': '/mnt', 'name': 'foo', 'readOnly': True, 'subPath': '/' }] }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'name': 'foo' }] } }) # Apply Pod Mutation Hook pod = self.pod_launcher._mutate_pod_backcompat(pod) sanitized_pod_post_mutation = api_client.sanitize_for_serialization( pod) self.assertEqual( sanitized_pod_post_mutation, { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'namespace': 'airflow-tests' }, 'spec': { 'containers': [{ 'args': [], 'command': ['foo'], 'env': [{ 'name': 'TEST_USER', 'value': 'ADMIN' }], 'envFrom': [], 'image': 'test-image', 'name': 'base', 'ports': [{ 'containerPort': 8080 }, { 'containerPort': 8081 }], 'volumeMounts': [{ 'mountPath': '/mnt', 'name': 'foo', 'readOnly': True, 'subPath': '/' }, { 'mountPath': '/opt/airflow/secrets/', 'name': 'airflow-secrets-mount', 'readOnly': True }] }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'name': 'foo' }, { 'name': 'airflow-secrets-mount', 'secret': { 'secretName': 'airflow-test-secrets' } }] } })
def test_pod_mutation_to_k8s_pod(self): with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK, "airflow_local_settings"): from airflow import settings settings.import_local_settings() # pylint: ignore from airflow.kubernetes.pod_launcher import PodLauncher self.mock_kube_client = Mock() self.pod_launcher = PodLauncher(kube_client=self.mock_kube_client) init_container = k8s.V1Container(name="init-container", volume_mounts=[ k8s.V1VolumeMount( mount_path="/tmp", name="init-secret") ]) pod = pod_generator.PodGenerator( image="foo", name="bar", namespace="baz", image_pull_policy="Never", init_containers=[init_container], cmds=["foo"], args=["/bin/sh", "-c", "touch /tmp/healthy"], tolerations=[{ 'effect': 'NoSchedule', 'key': 'static-pods', 'operator': 'Equal', 'value': 'true' }], volume_mounts=[{ "name": "foo", "mountPath": "/mnt", "subPath": "/", "readOnly": True }], security_context=k8s.V1PodSecurityContext(fs_group=0, run_as_user=1), volumes=[k8s.V1Volume(name="foo")]).gen_pod() sanitized_pod_pre_mutation = api_client.sanitize_for_serialization( pod) self.assertEqual( sanitized_pod_pre_mutation, { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': mock.ANY, 'namespace': 'baz' }, 'spec': { 'containers': [{ 'args': ['/bin/sh', '-c', 'touch /tmp/healthy'], 'command': ['foo'], 'env': [], 'envFrom': [], 'image': 'foo', 'imagePullPolicy': 'Never', 'name': 'base', 'ports': [], 'volumeMounts': [{ 'mountPath': '/mnt', 'name': 'foo', 'readOnly': True, 'subPath': '/' }] }], 'initContainers': [{ 'name': 'init-container', 'volumeMounts': [{ 'mountPath': '/tmp', 'name': 'init-secret' }] }], 'hostNetwork': False, 'imagePullSecrets': [], 'tolerations': [{ 'effect': 'NoSchedule', 'key': 'static-pods', 'operator': 'Equal', 'value': 'true' }], 'volumes': [{ 'name': 'foo' }], 'securityContext': { 'fsGroup': 0, 'runAsUser': 1 } } }, ) # Apply Pod Mutation Hook pod = self.pod_launcher._mutate_pod_backcompat(pod) sanitized_pod_post_mutation = api_client.sanitize_for_serialization( pod) self.assertEqual( sanitized_pod_post_mutation, { "apiVersion": "v1", "kind": "Pod", 'metadata': { 'labels': { 'test_label': 'test_value' }, 'name': mock.ANY, 'namespace': 'airflow-tests' }, 'spec': { 'affinity': { 'nodeAffinity': { 'requiredDuringSchedulingIgnoredDuringExecution': { 'nodeSelectorTerms': [{ 'matchExpressions': [{ 'key': 'test/dynamic-pods', 'operator': 'In', 'values': ['true'] }] }] } } }, 'containers': [{ 'args': ['/bin/sh', '-c', 'touch /tmp/healthy2'], 'command': ['foo'], 'env': [{ 'name': 'TEST_USER', 'value': 'ADMIN' }], 'image': 'my_image', 'imagePullPolicy': 'Never', 'name': 'base', 'ports': [{ 'containerPort': 8080 }, { 'containerPort': 8081 }], 'resources': { 'limits': { 'nvidia.com/gpu': '200G' }, 'requests': { 'cpu': '200Mi', 'memory': '2G' } }, 'volumeMounts': [{ 'mountPath': '/mnt', 'name': 'foo', 'readOnly': True, 'subPath': '/' }, { 'mountPath': '/opt/airflow/secrets/', 'name': 'airflow-secrets-mount', 'readOnly': True }] }], 'hostNetwork': False, 'imagePullSecrets': [], 'initContainers': [{ 'name': 'init-container', 'securityContext': { 'runAsGroup': 50000, 'runAsUser': 50000 }, 'volumeMounts': [{ 'mountPath': '/tmp', 'name': 'init-secret' }] }], 'tolerations': [{ 'effect': 'NoSchedule', 'key': 'static-pods', 'operator': 'Equal', 'value': 'true' }, { 'effect': 'NoSchedule', 'key': 'dynamic-pods', 'operator': 'Equal', 'value': 'true' }], 'volumes': [ { 'name': 'airflow-secrets-mount', 'secret': { 'secretName': 'airflow-test-secrets' } }, { 'name': 'bar' }, { 'name': 'foo' }, ], 'securityContext': { 'runAsUser': 1 } } })