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
}
}
})
