def build_gcsfs_controller_for_tests(
controller_cls,
fixture_path_prefix: str,
run_async: bool,
**kwargs,
) -> GcsfsDirectIngestController:
fake_fs = FakeDirectIngestGCSFileSystem()
def mock_build_fs():
return fake_fs
with patch(
'recidiviz.ingest.direct.controllers.'
'base_direct_ingest_controller.DirectIngestCloudTaskManagerImpl') \
as mock_task_factory_cls:
task_manager = FakeAsyncDirectIngestCloudTaskManager() \
if run_async else FakeSynchronousDirectIngestCloudTaskManager()
mock_task_factory_cls.return_value = task_manager
with patch.object(GcsfsFactory, 'build', new=mock_build_fs):
controller = controller_cls(
ingest_directory_path=f'{fixture_path_prefix}/fixtures',
storage_directory_path='storage/path',
**kwargs)
task_manager.set_controller(controller)
fake_fs.test_set_controller(controller)
return controller
def build_gcsfs_controller_for_tests(
controller_cls,
fixture_path_prefix: str,
run_async: bool,
fake_fs: Optional[FakeGCSFileSystem] = None,
can_start_ingest: bool = True,
**kwargs,
) -> GcsfsDirectIngestController:
"""Builds an instance of |controller_cls| for use in tests with several internal classes mocked properly. """
fake_fs = FakeGCSFileSystem()
def mock_build_fs():
return fake_fs
if 'TestGcsfsDirectIngestController' in controller_cls.__name__:
view_collector_cls: Type[BigQueryViewCollector] = \
FakeDirectIngestPreProcessedIngestViewCollector
else:
view_collector_cls = DirectIngestPreProcessedIngestViewCollector
with patch(
f'{BaseDirectIngestController.__module__}.DirectIngestCloudTaskManagerImpl'
) as mock_task_factory_cls:
with patch(
f'{GcsfsDirectIngestController.__module__}.BigQueryClientImpl'
) as mock_big_query_client_cls:
with patch(
f'{GcsfsDirectIngestController.__module__}.DirectIngestRawFileImportManager',
FakeDirectIngestRawFileImportManager):
with patch(
f'{GcsfsDirectIngestController.__module__}.DirectIngestPreProcessedIngestViewCollector',
view_collector_cls):
task_manager = FakeAsyncDirectIngestCloudTaskManager() \
if run_async else FakeSynchronousDirectIngestCloudTaskManager()
mock_task_factory_cls.return_value = task_manager
mock_big_query_client_cls.return_value = \
FakeDirectIngestBigQueryClient(project_id=metadata.project_id(), fs=fake_fs)
with patch.object(GcsfsFactory, 'build',
new=mock_build_fs):
controller = controller_cls(
ingest_directory_path=
f'{fixture_path_prefix}/fixtures',
storage_directory_path='storage/path',
**kwargs)
controller.csv_reader = TestSafeGcsCsvReader(fake_fs)
controller.raw_file_import_manager.csv_reader = controller.csv_reader
task_manager.set_controller(controller)
fake_fs.test_set_delegate(
DirectIngestFakeGCSFileSystemDelegate(
controller, can_start_ingest=can_start_ingest))
return controller
def build_controller_for_tests(controller_cls,
run_async: bool) -> BaseDirectIngestController:
if issubclass(controller_cls, GcsfsDirectIngestController):
raise ValueError(f"Controller class {controller_cls} is instance of "
f"GcsfsDirectIngestController - use "
f"build_gcsfs_controller_for_tests instead.")
with patch(
'recidiviz.ingest.direct.controllers.'
'base_direct_ingest_controller.DirectIngestCloudTaskManagerImpl') \
as mock_task_factory_cls:
task_manager = FakeAsyncDirectIngestCloudTaskManager() \
if run_async else FakeSynchronousDirectIngestCloudTaskManager()
mock_task_factory_cls.return_value = task_manager
return controller_cls()
def build_controller_for_tests(
controller_cls: Type[BaseDirectIngestController],
run_async: bool) -> BaseDirectIngestController:
"""Builds an instance of |controller_cls| for use in tests with several internal classes mocked properly. If
|controller_cls| is an instance of GcsfsDirectIngestController, use build_gcsfs_controller_for_tests() instead.
"""
if issubclass(controller_cls, GcsfsDirectIngestController):
raise ValueError(f"Controller class {controller_cls} is instance of "
f"GcsfsDirectIngestController - use "
f"build_gcsfs_controller_for_tests instead.")
with patch("recidiviz.ingest.direct.controllers."
"base_direct_ingest_controller.DirectIngestCloudTaskManagerImpl"
) as mock_task_factory_cls:
task_manager = (FakeAsyncDirectIngestCloudTaskManager() if run_async
else FakeSynchronousDirectIngestCloudTaskManager())
mock_task_factory_cls.return_value = task_manager
return controller_cls() # type: ignore[call-arg]
def build_gcsfs_controller_for_tests(
controller_cls: Type[CsvGcsfsDirectIngestController],
ingest_instance: DirectIngestInstance,
run_async: bool,
can_start_ingest: bool = True,
regions_module: ModuleType = fake_regions_module,
) -> BaseDirectIngestController:
"""Builds an instance of |controller_cls| for use in tests with several internal classes mocked properly. """
fake_fs = FakeGCSFileSystem()
def mock_build_fs() -> FakeGCSFileSystem:
return fake_fs
if "TestGcsfsDirectIngestController" in controller_cls.__name__:
view_collector_cls: Type[
BigQueryViewCollector] = FakeDirectIngestPreProcessedIngestViewCollector
else:
view_collector_cls = DirectIngestPreProcessedIngestViewCollector
with patch(
f"{BaseDirectIngestController.__module__}.DirectIngestCloudTaskManagerImpl"
) as mock_task_factory_cls:
with patch(
f"{BaseDirectIngestController.__module__}.BigQueryClientImpl"
) as mock_big_query_client_cls:
with patch(
f"{BaseDirectIngestController.__module__}.DirectIngestRawFileImportManager",
FakeDirectIngestRawFileImportManager,
):
with patch(
f"{BaseDirectIngestController.__module__}.DirectIngestPreProcessedIngestViewCollector",
view_collector_cls,
):
task_manager = (
FakeAsyncDirectIngestCloudTaskManager() if run_async
else FakeSynchronousDirectIngestCloudTaskManager())
mock_task_factory_cls.return_value = task_manager
mock_big_query_client_cls.return_value = (
FakeDirectIngestBigQueryClient(
project_id=metadata.project_id(),
fs=fake_fs,
region_code=controller_cls.region_code(),
))
with patch.object(GcsfsFactory, "build",
new=mock_build_fs):
with patch.object(
direct_ingest_raw_table_migration_collector,
"regions",
new=regions_module,
):
controller = controller_cls(
ingest_bucket_path=
gcsfs_direct_ingest_bucket_for_region(
region_code=controller_cls.region_code(),
system_level=SystemLevel.for_region_code(
controller_cls.region_code(),
is_direct_ingest=True,
),
ingest_instance=ingest_instance,
project_id="recidiviz-xxx",
))
controller.csv_reader = GcsfsCsvReader(fake_fs)
controller.raw_file_import_manager.csv_reader = (
controller.csv_reader)
task_manager.set_controller(controller)
fake_fs.test_set_delegate(
DirectIngestFakeGCSFileSystemDelegate(
controller,
can_start_ingest=can_start_ingest))
return controller