def test_proxy_ssl_with_verification(self, ssl_mock, http_response_mock,
create_connection_mock):
type(http_response_mock.return_value).status = mock.PropertyMock(
return_value=200)
conn = AWSAuthConnection('mockservice.s3.amazonaws.com',
aws_access_key_id='access_key',
aws_secret_access_key='secret',
suppress_consec_slashes=False,
proxy_port=80)
conn.https_validate_certificates = True
dummy_cert = {
'subjectAltName':
(('DNS', 's3.amazonaws.com'), ('DNS', '*.s3.amazonaws.com')),
}
mock_sock = mock.Mock()
create_connection_mock.return_value = mock_sock
mock_sslSock = mock.Mock()
mock_sslSock.getpeercert.return_value = dummy_cert
mock_context = mock.Mock()
mock_context.wrap_socket.return_value = mock_sslSock
ssl_mock.create_default_context.return_value = mock_context
# Attempt to call proxy_ssl and make sure it works
conn.proxy_ssl('mockservice.s3.amazonaws.com', 80)
mock_sslSock.getpeercert.assert_called_once_with()
mock_context.wrap_socket.assert_called_once_with(
mock_sock, server_hostname='mockservice.s3.amazonaws.com')
def test_synthesize_text_green_path(self, mock_text_to_speech_hook,
mock_gcp_hook):
mocked_response = mock.Mock()
type(mocked_response).audio_content = mock.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()})
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=TARGET_BUCKET_NAME,
object=TARGET_FILENAME,
filename=mock.ANY)
def test_should_copy_single_file(self, mock_named_temporary_file,
mock_gdrive, mock_gcs_hook):
type(mock_named_temporary_file.return_value.__enter__.return_value
).name = mock.PropertyMock(side_effect=["TMP1"])
task = GcsToGDriveOperator(
task_id="copy_single_file",
source_bucket="data",
source_object="sales/sales-2017/january.avro",
destination_object="copied_sales/2017/january-backup.avro",
)
task.execute(mock.MagicMock())
mock_gcs_hook.assert_has_calls([
mock.call(delegate_to=None,
google_cloud_storage_conn_id="google_cloud_default"),
mock.call().download(bucket="data",
filename="TMP1",
object="sales/sales-2017/january.avro"),
])
mock_gdrive.assert_has_calls([
mock.call(delegate_to=None, gcp_conn_id="google_cloud_default"),
mock.call().upload_file(
local_location="TMP1",
remote_location="copied_sales/2017/january-backup.avro"),
])
def test_submit_run(self, mock_requests):
mock_requests.codes.ok = 200
mock_requests.post.return_value.json.return_value = {'run_id': '1'}
status_code_mock = mock.PropertyMock(return_value=200)
type(mock_requests.post.return_value).status_code = status_code_mock
json = {'notebook_task': NOTEBOOK_TASK, 'new_cluster': NEW_CLUSTER}
run_id = self.hook.submit_run(json)
self.assertEqual(run_id, '1')
args = mock_requests.post.call_args
kwargs = args[1]
self.assertEqual(kwargs['auth'].token, TOKEN)
def test_terminate_cluster(self, mock_requests):
mock_requests.codes.ok = 200
mock_requests.post.return_value.json.return_value = {}
status_code_mock = mock.PropertyMock(return_value=200)
type(mock_requests.post.return_value).status_code = status_code_mock
self.hook.terminate_cluster({"cluster_id": CLUSTER_ID})
mock_requests.post.assert_called_once_with(
terminate_cluster_endpoint(HOST),
json={'cluster_id': CLUSTER_ID},
auth=(LOGIN, PASSWORD),
headers=USER_AGENT_HEADER,
timeout=self.hook.timeout_seconds)
def test_proxy_ssl(self, ssl_mock, http_response_mock,
create_connection_mock):
type(http_response_mock.return_value).status = mock.PropertyMock(
return_value=200)
conn = AWSAuthConnection('mockservice.cc-zone-1.amazonaws.com',
aws_access_key_id='access_key',
aws_secret_access_key='secret',
suppress_consec_slashes=False,
proxy_port=80)
conn.https_validate_certificates = False
# Attempt to call proxy_ssl and make sure it works
conn.proxy_ssl('mockservice.cc-zone-1.amazonaws.com', 80)
def test_should_move_files(self, mock_named_temporary_file, mock_gdrive,
mock_gcs_hook):
type(mock_named_temporary_file.return_value.__enter__.return_value
).name = mock.PropertyMock(side_effect=["TMP1", "TMP2", "TMP3"])
mock_gcs_hook.return_value.list.return_value = [
"sales/A.avro", "sales/B.avro", "sales/C.avro"
]
task = GcsToGDriveOperator(
task_id="move_files",
source_bucket="data",
source_object="sales/sales-2017/*.avro",
move_object=True,
)
task.execute(mock.MagicMock())
mock_gcs_hook.assert_has_calls([
mock.call(delegate_to=None,
google_cloud_storage_conn_id="google_cloud_default"),
mock.call().list("data",
delimiter=".avro",
prefix="sales/sales-2017/"),
mock.call().download(bucket="data",
filename="TMP1",
object="sales/A.avro"),
mock.call().delete("data", "sales/A.avro"),
mock.call().download(bucket="data",
filename="TMP2",
object="sales/B.avro"),
mock.call().delete("data", "sales/B.avro"),
mock.call().download(bucket="data",
filename="TMP3",
object="sales/C.avro"),
mock.call().delete("data", "sales/C.avro"),
])
mock_gdrive.assert_has_calls([
mock.call(delegate_to=None, gcp_conn_id="google_cloud_default"),
mock.call().upload_file(local_location="TMP1",
remote_location="sales/A.avro"),
mock.call().upload_file(local_location="TMP2",
remote_location="sales/B.avro"),
mock.call().upload_file(local_location="TMP3",
remote_location="sales/C.avro"),
])
def test_run_now(self, mock_requests):
mock_requests.codes.ok = 200
mock_requests.post.return_value.json.return_value = {'run_id': '1'}
status_code_mock = mock.PropertyMock(return_value=200)
type(mock_requests.post.return_value).status_code = status_code_mock
json = {
'notebook_params': NOTEBOOK_PARAMS,
'jar_params': JAR_PARAMS,
'job_id': JOB_ID
}
run_id = self.hook.run_now(json)
self.assertEqual(run_id, '1')
mock_requests.post.assert_called_once_with(
run_now_endpoint(HOST),
json={
'notebook_params': NOTEBOOK_PARAMS,
'jar_params': JAR_PARAMS,
'job_id': JOB_ID
},
auth=(LOGIN, PASSWORD),
headers=USER_AGENT_HEADER,
timeout=self.hook.timeout_seconds)
