def test_exec_success_json(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test successful run of execute function for JSON""" op = MySqlToGoogleCloudStorageOperator(task_id=TASK_ID, mysql_conn_id=MYSQL_CONN_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME) mysql_hook_mock = mysql_hook_mock_class.return_value mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION gcs_hook_mock = gcs_hook_mock_class.return_value def _assert_upload(bucket, obj, tmp_filename, mime_type=None): self.assertEqual(BUCKET, bucket) self.assertEqual(JSON_FILENAME.format(0), obj) self.assertEqual('application/json', mime_type) with open(tmp_filename, 'rb') as f: self.assertEqual(b''.join(NDJSON_LINES), f.read()) gcs_hook_mock.upload.side_effect = _assert_upload op.execute(None) mysql_hook_mock_class.assert_called_once_with( mysql_conn_id=MYSQL_CONN_ID) mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with( SQL)
def test_exec_success_csv_with_delimiter(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test successful run of execute function for CSV with a field delimiter""" op = MySqlToGoogleCloudStorageOperator(task_id=TASK_ID, mysql_conn_id=MYSQL_CONN_ID, sql=SQL, export_format='csv', field_delimiter='|', bucket=BUCKET, filename=CSV_FILENAME) mysql_hook_mock = mysql_hook_mock_class.return_value mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION gcs_hook_mock = gcs_hook_mock_class.return_value def _assert_upload(bucket, obj, tmp_filename, mime_type=None): self.assertEqual(BUCKET, bucket) self.assertEqual(CSV_FILENAME.format(0), obj) self.assertEqual('text/csv', mime_type) with open(tmp_filename, 'rb') as f: self.assertEqual(b''.join(CSV_LINES_PIPE_DELIMITED), f.read()) gcs_hook_mock.upload.side_effect = _assert_upload op.execute(None) mysql_hook_mock_class.assert_called_once_with( mysql_conn_id=MYSQL_CONN_ID) mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with( SQL)
def test_file_splitting(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test that ndjson is split by approx_max_file_size_bytes param.""" mysql_hook_mock = mysql_hook_mock_class.return_value mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION gcs_hook_mock = gcs_hook_mock_class.return_value expected_upload = { JSON_FILENAME.format(0): b''.join(NDJSON_LINES[:2]), JSON_FILENAME.format(1): NDJSON_LINES[2], } def _assert_upload(bucket, obj, tmp_filename, mime_type=None): self.assertEqual(BUCKET, bucket) self.assertEqual('application/json', mime_type) with open(tmp_filename, 'rb') as f: self.assertEqual(expected_upload[obj], f.read()) gcs_hook_mock.upload.side_effect = _assert_upload op = MySqlToGoogleCloudStorageOperator( task_id=TASK_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME, approx_max_file_size_bytes=len( expected_upload[JSON_FILENAME.format(0)])) op.execute(None)
def test_schema_file(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test writing schema files.""" mysql_hook_mock = mysql_hook_mock_class.return_value mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION gcs_hook_mock = gcs_hook_mock_class.return_value def _assert_upload(bucket, obj, tmp_filename, mime_type): if obj == SCHEMA_FILENAME: with open(tmp_filename, 'rb') as f: self.assertEqual(b''.join(SCHEMA_JSON), f.read()) gcs_hook_mock.upload.side_effect = _assert_upload op = MySqlToGoogleCloudStorageOperator( task_id=TASK_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME, schema_filename=SCHEMA_FILENAME) op.execute(None) # once for the file and once for the schema self.assertEqual(2, gcs_hook_mock.upload.call_count)
def test_exec_success_json(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test successful run of execute function for JSON""" op = MySqlToGoogleCloudStorageOperator( task_id=TASK_ID, mysql_conn_id=MYSQL_CONN_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME) mysql_hook_mock = mysql_hook_mock_class.return_value mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION gcs_hook_mock = gcs_hook_mock_class.return_value def _assert_upload(bucket, obj, tmp_filename, mime_type=None): self.assertEqual(BUCKET, bucket) self.assertEqual(JSON_FILENAME.format(0), obj) self.assertEqual('application/json', mime_type) with open(tmp_filename, 'rb') as f: self.assertEqual(b''.join(NDJSON_LINES), f.read()) gcs_hook_mock.upload.side_effect = _assert_upload op.execute(None) mysql_hook_mock_class.assert_called_once_with(mysql_conn_id=MYSQL_CONN_ID) mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL)
def test_file_splitting(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test that ndjson is split by approx_max_file_size_bytes param.""" mysql_hook_mock = mysql_hook_mock_class.return_value mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION gcs_hook_mock = gcs_hook_mock_class.return_value expected_upload = { JSON_FILENAME.format(0): b''.join(NDJSON_LINES[:2]), JSON_FILENAME.format(1): NDJSON_LINES[2], } def _assert_upload(bucket, obj, tmp_filename, mime_type=None): self.assertEqual(BUCKET, bucket) self.assertEqual('application/json', mime_type) with open(tmp_filename, 'rb') as f: self.assertEqual(expected_upload[obj], f.read()) gcs_hook_mock.upload.side_effect = _assert_upload op = MySqlToGoogleCloudStorageOperator( task_id=TASK_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME, approx_max_file_size_bytes=len(expected_upload[JSON_FILENAME.format(0)])) op.execute(None)
def test_exec_success_csv_with_delimiter(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test successful run of execute function for CSV with a field delimiter""" op = MySqlToGoogleCloudStorageOperator( task_id=TASK_ID, mysql_conn_id=MYSQL_CONN_ID, sql=SQL, export_format='csv', field_delimiter='|', bucket=BUCKET, filename=CSV_FILENAME) mysql_hook_mock = mysql_hook_mock_class.return_value mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION gcs_hook_mock = gcs_hook_mock_class.return_value def _assert_upload(bucket, obj, tmp_filename, mime_type=None): self.assertEqual(BUCKET, bucket) self.assertEqual(CSV_FILENAME.format(0), obj) self.assertEqual('text/csv', mime_type) with open(tmp_filename, 'rb') as f: self.assertEqual(b''.join(CSV_LINES_PIPE_DELIMITED), f.read()) gcs_hook_mock.upload.side_effect = _assert_upload op.execute(None) mysql_hook_mock_class.assert_called_once_with(mysql_conn_id=MYSQL_CONN_ID) mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL)
def test_schema_file(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test writing schema files.""" mysql_hook_mock = mysql_hook_mock_class.return_value mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION gcs_hook_mock = gcs_hook_mock_class.return_value def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip): # pylint: disable=unused-argument if obj == SCHEMA_FILENAME: self.assertFalse(gzip) with open(tmp_filename, 'rb') as file: self.assertEqual(b''.join(SCHEMA_JSON), file.read()) gcs_hook_mock.upload.side_effect = _assert_upload op = MySqlToGoogleCloudStorageOperator( task_id=TASK_ID, sql=SQL, bucket=BUCKET, filename=JSON_FILENAME, schema_filename=SCHEMA_FILENAME) op.execute(None) # once for the file and once for the schema self.assertEqual(2, gcs_hook_mock.upload.call_count)
def test_exec_success_csv_ensure_utc(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test successful run of execute function for CSV""" op = MySqlToGoogleCloudStorageOperator( task_id=TASK_ID, mysql_conn_id=MYSQL_CONN_ID, sql=SQL, export_format='CSV', bucket=BUCKET, filename=CSV_FILENAME, ensure_utc=True) mysql_hook_mock = mysql_hook_mock_class.return_value mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION gcs_hook_mock = gcs_hook_mock_class.return_value def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False): self.assertEqual(BUCKET, bucket) self.assertEqual(CSV_FILENAME.format(0), obj) self.assertEqual('text/csv', mime_type) self.assertFalse(gzip) with open(tmp_filename, 'rb') as file: self.assertEqual(b''.join(CSV_LINES), file.read()) gcs_hook_mock.upload.side_effect = _assert_upload op.execute(None) mysql_hook_mock_class.assert_called_once_with(mysql_conn_id=MYSQL_CONN_ID) mysql_hook_mock.get_conn().cursor().execute.assert_has_calls([mock.call(TZ_QUERY), mock.call(SQL)])
def test_exec_success_csv(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test the execute function in case where the run is successful.""" op = MySqlToGoogleCloudStorageOperator( task_id=TASK_ID, mysql_conn_id=MYSQL_CONN_ID, sql=SQL, export_format={'file_format': 'csv', 'csv_dialect': 'excel'}, bucket=BUCKET, filename=FILENAME) mysql_hook_mock = mysql_hook_mock_class.return_value mysql_hook_mock.get_conn().cursor().__iter__.return_value = iter(ROWS) mysql_hook_mock.get_conn().cursor().description = CURSOR_DESCRIPTION gcs_hook_mock = gcs_hook_mock_class.return_value def _assert_upload(bucket, obj, tmp_filename, content_type): self.assertEqual(BUCKET, bucket) self.assertEqual(FILENAME.format(0), obj) self.assertEqual('application/csv', content_type) with open(tmp_filename, 'rb') as f: self.assertEqual(b''.join(CSV_LINES), f.read()) gcs_hook_mock.upload.side_effect = _assert_upload op.execute(None) mysql_hook_mock_class.assert_called_once_with(mysql_conn_id=MYSQL_CONN_ID) mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL)