def test_raw_data_import(self, mock_supported, mock_region,
mock_environment):
mock_supported.return_value = ['us_xx']
region_code = 'us_xx'
mock_environment.return_value = 'staging'
mock_controller = create_autospec(GcsfsDirectIngestController)
mock_region.return_value = fake_region(region_code=region_code,
environment='staging',
ingestor=mock_controller)
import_args = GcsfsRawDataBQImportArgs(
raw_data_file_path=GcsfsFilePath.from_absolute_path(
to_normalized_unprocessed_file_path(
'bucket/raw_data_path.csv',
file_type=GcsfsDirectIngestFileType.RAW_DATA)))
request_args = {
'region': region_code,
}
body = {
'cloud_task_args': import_args.to_serializable(),
'args_type': 'GcsfsRawDataBQImportArgs',
}
body_encoded = json.dumps(body).encode()
headers = {'X-Appengine-Cron': 'test-cron'}
response = self.client.post('/raw_data_import',
query_string=request_args,
headers=headers,
data=body_encoded)
self.assertEqual(200, response.status_code)
mock_controller.do_raw_data_import.assert_called_with(import_args)
def test_create_direct_ingest_raw_data_import_task(
self, mock_client: mock.MagicMock,
mock_uuid: mock.MagicMock) -> None:
# Arrange
raw_data_path = GcsfsFilePath.from_absolute_path(
to_normalized_unprocessed_file_path(
"bucket/raw_data_path.csv",
file_type=GcsfsDirectIngestFileType.RAW_DATA,
))
import_args = GcsfsRawDataBQImportArgs(
raw_data_file_path=raw_data_path)
body = {
"cloud_task_args": import_args.to_serializable(),
"args_type": "GcsfsRawDataBQImportArgs",
}
body_encoded = json.dumps(body).encode()
uuid = "random-uuid"
mock_uuid.uuid4.return_value = uuid
date = "2019-07-20"
queue_path = f"{DIRECT_INGEST_BQ_IMPORT_EXPORT_QUEUE_V2}-path"
task_name = DIRECT_INGEST_BQ_IMPORT_EXPORT_QUEUE_V2 + "/{}-{}-{}".format(
_REGION.region_code, date, uuid)
url_params = {
"region": _REGION.region_code,
"file_path": raw_data_path.abs_path(),
}
task = tasks_v2.types.task_pb2.Task(
name=task_name,
app_engine_http_request={
"http_method": "POST",
"relative_uri":
f"/direct/raw_data_import?{urlencode(url_params)}",
"body": body_encoded,
},
)
mock_client.return_value.task_path.return_value = task_name
mock_client.return_value.queue_path.return_value = queue_path
# Act
DirectIngestCloudTaskManagerImpl(
).create_direct_ingest_raw_data_import_task(
_REGION, DirectIngestInstance.PRIMARY, import_args)
# Assert
mock_client.return_value.queue_path.assert_called_with(
self.mock_project_id, QUEUES_REGION,
DIRECT_INGEST_BQ_IMPORT_EXPORT_QUEUE_V2)
mock_client.return_value.create_task.assert_called_with(
parent=queue_path, task=task)
def test_create_direct_ingest_raw_data_import_task(self, mock_client,
mock_uuid):
# Arrange
project_id = 'recidiviz-456'
import_args = GcsfsRawDataBQImportArgs(
raw_data_file_path=GcsfsFilePath.from_absolute_path(
to_normalized_unprocessed_file_path(
'bucket/raw_data_path.csv',
file_type=GcsfsDirectIngestFileType.RAW_DATA)))
body = {
'cloud_task_args': import_args.to_serializable(),
'args_type': 'GcsfsRawDataBQImportArgs'
}
body_encoded = json.dumps(body).encode()
uuid = 'random-uuid'
mock_uuid.uuid4.return_value = uuid
date = '2019-07-20'
queue_path = _REGION.shared_queue + '-path'
task_name = DIRECT_INGEST_BQ_IMPORT_EXPORT_QUEUE_V2 + '/{}-{}-{}'.format(
_REGION.region_code, date, uuid)
task = tasks_v2.types.task_pb2.Task(
name=task_name,
app_engine_http_request={
'http_method': 'POST',
'relative_uri':
f'/direct/raw_data_import?region={_REGION.region_code}',
'body': body_encoded
})
mock_client.return_value.task_path.return_value = task_name
mock_client.return_value.queue_path.return_value = queue_path
# Act
DirectIngestCloudTaskManagerImpl(
project_id=project_id).create_direct_ingest_raw_data_import_task(
_REGION, import_args)
# Assert
mock_client.return_value.queue_path.assert_called_with(
project_id, QUEUES_REGION, DIRECT_INGEST_BQ_IMPORT_EXPORT_QUEUE_V2)
mock_client.return_value.create_task.assert_called_with(
parent=queue_path, task=task)
def test_raw_data_import(
self,
mock_supported: mock.MagicMock,
mock_region: mock.MagicMock,
mock_environment: mock.MagicMock,
) -> None:
mock_supported.return_value = ["us_xx"]
region_code = "us_xx"
mock_environment.return_value = "staging"
mock_controller = create_autospec(GcsfsDirectIngestController)
mock_region.return_value = fake_region(region_code=region_code,
environment="staging",
ingestor=mock_controller)
import_args = GcsfsRawDataBQImportArgs(
raw_data_file_path=GcsfsFilePath.from_absolute_path(
to_normalized_unprocessed_file_path(
"bucket/raw_data_path.csv",
file_type=GcsfsDirectIngestFileType.RAW_DATA,
)))
request_args = {
"region": region_code,
}
body = {
"cloud_task_args": import_args.to_serializable(),
"args_type": "GcsfsRawDataBQImportArgs",
}
body_encoded = json.dumps(body).encode()
headers = {"X-Appengine-Cron": "test-cron"}
response = self.client.post(
"/raw_data_import",
query_string=request_args,
headers=headers,
data=body_encoded,
)
self.assertEqual(200, response.status_code)
mock_controller.do_raw_data_import.assert_called_with(import_args)