class ViewCollectionExportManagerTest(unittest.TestCase): """Tests for view_export_manager.py.""" def setUp(self) -> None: self.app = Flask(__name__) self.app.register_blueprint(export_blueprint) self.app.config["TESTING"] = True self.headers: Dict[str, Dict[Any, Any]] = { "x-goog-iap-jwt-assertion": {} } self.client = self.app.test_client() self.mock_cloud_task_client_patcher = mock.patch( "google.cloud.tasks_v2.CloudTasksClient") self.mock_cloud_task_client_patcher.start() self.mock_uuid_patcher = mock.patch( f"{CLOUD_TASK_MANAGER_PACKAGE_NAME}.uuid") self.mock_uuid = self.mock_uuid_patcher.start() with self.app.test_request_context(): self.metric_view_data_export_url = flask.url_for( "export.metric_view_data_export") self.create_metric_view_data_export_tasks_url = flask.url_for( "export.create_metric_view_data_export_tasks") self.mock_state_code = "US_XX" self.mock_project_id = "test-project" self.mock_dataset_id = "base_dataset" self.mock_dataset = bigquery.dataset.DatasetReference( self.mock_project_id, self.mock_dataset_id) self.metadata_patcher = mock.patch( "recidiviz.utils.metadata.project_id") self.mock_project_id_fn = self.metadata_patcher.start() self.mock_project_id_fn.return_value = self.mock_project_id self.client_patcher = mock.patch( "recidiviz.metrics.export.view_export_manager.BigQueryClientImpl") self.mock_client = self.client_patcher.start().return_value self.mock_client.dataset_ref_for_id.return_value = self.mock_dataset self.mock_view_builder = SimpleBigQueryViewBuilder( dataset_id=self.mock_dataset.dataset_id, view_id="test_view", description="test_view description", view_query_template="SELECT NULL LIMIT 0", ) self.mock_metric_view_builder = MetricBigQueryViewBuilder( dataset_id=self.mock_dataset.dataset_id, view_id="test_view", description="test_view description", view_query_template="SELECT NULL LIMIT 0", dimensions=tuple(), ) self.view_builders_for_dataset = [ self.mock_view_builder, self.mock_metric_view_builder, ] self.output_uri_template_for_dataset = { "dataset_id": "gs://{project_id}-dataset-location/subdirectory", } self.views_to_update = { self.mock_dataset_id: self.view_builders_for_dataset } self.mock_export_name = "MOCK_EXPORT_NAME" self.mock_big_query_view_namespace = BigQueryViewNamespace.STATE self.metric_dataset_export_configs_index = { "EXPORT": ExportViewCollectionConfig( view_builders_to_export=[self.mock_view_builder], output_directory_uri_template= "gs://{project_id}-dataset-location/subdirectory", export_name="EXPORT", bq_view_namespace=self.mock_big_query_view_namespace, ), "OTHER_EXPORT": ExportViewCollectionConfig( view_builders_to_export=[self.mock_metric_view_builder], output_directory_uri_template= "gs://{project_id}-dataset-location/subdirectory", export_name="OTHER_EXPORT", bq_view_namespace=self.mock_big_query_view_namespace, ), self.mock_export_name: ExportViewCollectionConfig( view_builders_to_export=self.view_builders_for_dataset, output_directory_uri_template= "gs://{project_id}-dataset-location/subdirectory", export_name=self.mock_export_name, bq_view_namespace=self.mock_big_query_view_namespace, ), } export_config_values = { "OUTPUT_DIRECTORY_URI_TEMPLATE_FOR_DATASET_EXPORT": self.output_uri_template_for_dataset, "VIEW_COLLECTION_EXPORT_INDEX": self.metric_dataset_export_configs_index, } self.export_config_patcher = mock.patch( # type: ignore[call-overload] "recidiviz.metrics.export.view_export_manager.export_config", **export_config_values, ) self.mock_export_config = self.export_config_patcher.start() self.gcs_factory_patcher = mock.patch( "recidiviz.metrics.export.view_export_manager.GcsfsFactory.build") self.gcs_factory_patcher.start().return_value = FakeGCSFileSystem() def tearDown(self) -> None: self.client_patcher.stop() self.export_config_patcher.stop() self.metadata_patcher.stop() self.gcs_factory_patcher.stop() self.mock_uuid_patcher.stop() self.mock_cloud_task_client_patcher.stop() @mock.patch("recidiviz.utils.environment.get_gcp_environment") def test_get_configs_for_export_name( self, mock_environment: mock.MagicMock) -> None: """Tests get_configs_for_export_name function to ensure that export names correctly match""" mock_environment.return_value = "production" export_configs_for_filter = view_export_manager.get_configs_for_export_name( export_name=self.mock_export_name, state_code=self.mock_state_code, project_id=self.mock_project_id, ) view = self.mock_view_builder.build() metric_view = self.mock_metric_view_builder.build() expected_view_config_list = [ ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=view, view_filter_clause= f" WHERE state_code = '{self.mock_state_code}'", intermediate_table_name= f"{view.view_id}_table_{self.mock_state_code}", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory/{state_code}" .format( project_id=self.mock_project_id, state_code=self.mock_state_code, )), export_output_formats=[ExportOutputFormatType.JSON], ), ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=metric_view, view_filter_clause= f" WHERE state_code = '{self.mock_state_code}'", intermediate_table_name= f"{view.view_id}_table_{self.mock_state_code}", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory/{state_code}" .format( project_id=self.mock_project_id, state_code=self.mock_state_code, )), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ), ] self.assertEqual(expected_view_config_list, export_configs_for_filter) # Test for case insensitivity export_configs_for_filter = view_export_manager.get_configs_for_export_name( export_name=self.mock_export_name.lower(), state_code=self.mock_state_code.lower(), project_id=self.mock_project_id, ) self.assertEqual(expected_view_config_list, export_configs_for_filter) @mock.patch("recidiviz.utils.environment.get_gcp_environment") def test_get_configs_for_export_name_state_agnostic( self, mock_environment: mock.MagicMock) -> None: """Tests get_configs_for_export_name function to ensure that export names correctly match""" mock_environment.return_value = "production" export_configs_for_filter = view_export_manager.get_configs_for_export_name( export_name=self.mock_export_name, project_id=self.mock_project_id) view = self.mock_view_builder.build() metric_view = self.mock_metric_view_builder.build() expected_view_config_list = [ ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=view, view_filter_clause=None, intermediate_table_name=f"{view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory".format( project_id=self.mock_project_id, )), export_output_formats=[ExportOutputFormatType.JSON], ), ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=metric_view, view_filter_clause=None, intermediate_table_name=f"{view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory".format( project_id=self.mock_project_id, )), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ), ] self.assertEqual(expected_view_config_list, export_configs_for_filter) # Test for case insensitivity export_configs_for_filter = view_export_manager.get_configs_for_export_name( export_name=self.mock_export_name.lower(), project_id=self.mock_project_id) self.assertEqual(expected_view_config_list, export_configs_for_filter) @mock.patch("recidiviz.big_query.view_update_manager.rematerialize_views") @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage( self, mock_view_exporter: Mock, mock_view_update_manager_rematerialize: Mock) -> None: """Tests the table is created from the view and then extracted.""" view_export_manager.export_view_data_to_cloud_storage( self.mock_export_name, self.mock_state_code, mock_view_exporter) view = self.mock_view_builder.build() metric_view = self.mock_metric_view_builder.build() expected_view_config_list_1 = [ ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory/{state_code}" .format( project_id=self.mock_project_id, state_code="US_XX", )), export_output_formats=[ExportOutputFormatType.JSON], ) ] expected_view_config_list_2 = [ ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=metric_view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory/{state_code}" .format( project_id=self.mock_project_id, state_code="US_XX", )), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ) ] mock_view_update_manager_rematerialize.assert_called() mock_view_exporter.export_and_validate.assert_has_calls( [ mock.call([]), # CSV export mock.call([]), mock.call([ expected_view_config_list_1[0]. pointed_to_staging_subdirectory(), expected_view_config_list_2[0]. pointed_to_staging_subdirectory(), ]), # JSON exports mock.call([ expected_view_config_list_2[0]. pointed_to_staging_subdirectory() ]), # METRIC export ("OTHER_EXPORT") ], any_order=True, ) @mock.patch( "recidiviz.big_query.view_update_manager.create_managed_dataset_and_deploy_views_for_view_builders" ) @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_raise_exception_no_export_matched( self, mock_view_exporter: Mock, mock_view_update_manager_rematerialize: Mock) -> None: # pylint: disable=unused-argument """Tests the table is created from the view and then extracted.""" self.mock_export_config.NAMESPACE_TO_UPDATE_FOR_EXPORT_FILTER = { "US_YY": "NAMESPACE" } with self.assertRaises(ValueError) as e: view_export_manager.export_view_data_to_cloud_storage( export_job_name="JOBZZZ", override_view_exporter=mock_view_exporter) self.assertEqual( str(e.exception), "Export filter did not match any export configs:", " JOBZZZ", ) @mock.patch("recidiviz.big_query.view_update_manager.rematerialize_views") @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage_state_agnostic( self, mock_view_exporter: Mock, mock_view_update_manager_rematerialize: Mock) -> None: """Tests the table is created from the view and then extracted, where the export is not state-specific.""" state_agnostic_dataset_export_configs = { self.mock_export_name: ExportViewCollectionConfig( view_builders_to_export=self.view_builders_for_dataset, output_directory_uri_template= "gs://{project_id}-bucket-without-state-codes", export_name=self.mock_export_name, bq_view_namespace=self.mock_big_query_view_namespace, ), } self.mock_export_config.VIEW_COLLECTION_EXPORT_INDEX = ( state_agnostic_dataset_export_configs) view_export_manager.export_view_data_to_cloud_storage( export_job_name=self.mock_export_name, override_view_exporter=mock_view_exporter, ) view = self.mock_view_builder.build() metric_view = self.mock_metric_view_builder.build() view_export_configs = [ ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=view, view_filter_clause=None, intermediate_table_name=f"{view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-bucket-without-state-codes".format( project_id=self.mock_project_id, )), export_output_formats=[ExportOutputFormatType.JSON], ), ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=metric_view, view_filter_clause=None, intermediate_table_name=f"{view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-bucket-without-state-codes".format( project_id=self.mock_project_id, )), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ), ] mock_view_update_manager_rematerialize.assert_called() mock_view_exporter.export_and_validate.assert_has_calls( [ mock.call([]), # CSV export mock.call([ view_export_configs[1].pointed_to_staging_subdirectory() ]), # JSON export mock.call([ conf.pointed_to_staging_subdirectory() for conf in view_export_configs ]), # METRIC export ], any_order=True, ) @mock.patch("recidiviz.big_query.view_update_manager.rematerialize_views") @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage_value_error( self, mock_view_exporter: Mock, mock_view_update_manager_rematerialize: Mock) -> None: """Tests the table is created from the view and then extracted.""" mock_view_exporter.export_and_validate.side_effect = ValueError with self.assertRaises(ValueError): view_export_manager.export_view_data_to_cloud_storage( self.mock_export_name, override_view_exporter=mock_view_exporter) # Just the metric export is attempted and then the raise stops subsequent checks from happening mock_view_update_manager_rematerialize.assert_called_once() @mock.patch("recidiviz.big_query.view_update_manager.rematerialize_views") @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage_validation_error( self, mock_view_exporter: Mock, mock_view_update_manager_rematerialize: Mock) -> None: """Tests the table is created from the view and then extracted.""" mock_view_exporter.export_and_validate.side_effect = ViewExportValidationError # Should not throw view_export_manager.export_view_data_to_cloud_storage( self.mock_export_name, override_view_exporter=mock_view_exporter) # Just the metric export is attempted and then the raise stops subsequent checks from happening mock_view_update_manager_rematerialize.assert_called_once() @mock.patch("recidiviz.metrics.export.view_export_manager.deployed_views") @mock.patch("recidiviz.big_query.view_update_manager.rematerialize_views") @mock.patch( "recidiviz.big_query.view_update_manager.create_managed_dataset_and_deploy_views_for_view_builders" ) @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage_update_all_views( self, mock_view_exporter: Mock, mock_view_update_manager_deploy: Mock, mock_view_update_manager_rematerialize: Mock, mock_deployed_views: Mock, ) -> None: """Tests that all views in the namespace are updated before the export when the export name is in export_config.NAMESPACES_REQUIRING_FULL_UPDATE.""" self.mock_export_config.NAMESPACES_REQUIRING_FULL_UPDATE = [ self.mock_big_query_view_namespace ] mock_deployed_views.DEPLOYED_VIEW_BUILDERS_BY_NAMESPACE = { self.mock_big_query_view_namespace: self.view_builders_for_dataset } view_export_manager.export_view_data_to_cloud_storage( self.mock_export_name, override_view_exporter=mock_view_exporter) mock_view_update_manager_deploy.assert_called_with( view_source_table_datasets=VIEW_SOURCE_TABLE_DATASETS, view_builders_to_update=self.view_builders_for_dataset, ) mock_view_update_manager_rematerialize.assert_called_once() @mock.patch("recidiviz.metrics.export.view_export_manager.deployed_views") @mock.patch("recidiviz.big_query.view_update_manager.rematerialize_views") @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage_update_materialized_views_only( self, mock_view_exporter: Mock, mock_view_update_manager_rematerialize: Mock, mock_deployed_views: Mock, ) -> None: """Tests that only materialized views in the namespace are updated before the export when the export name is not in export_config.NAMESPACES_REQUIRING_FULL_UPDATE.""" self.mock_export_config.NAMESPACES_REQUIRING_FULL_UPDATE = [ "OTHER_NAMESPACE" ] mock_deployed_views.DEPLOYED_VIEW_BUILDERS_BY_NAMESPACE = { self.mock_big_query_view_namespace: self.view_builders_for_dataset } view_export_manager.export_view_data_to_cloud_storage( self.mock_export_name, override_view_exporter=mock_view_exporter) mock_view_update_manager_rematerialize.assert_called_with( view_source_table_datasets=VIEW_SOURCE_TABLE_DATASETS, all_view_builders=DEPLOYED_VIEW_BUILDERS, views_to_update=[ view.build() for view in self.view_builders_for_dataset ], ) @mock.patch( "recidiviz.metrics.export.view_export_manager.export_view_data_to_cloud_storage" ) def test_metric_view_data_export_valid_request( self, mock_export_view_data_to_cloud_storage: Mock) -> None: with self.app.test_request_context(): mock_export_view_data_to_cloud_storage.return_value = None response = self.client.get( self.metric_view_data_export_url, headers=self.headers, query_string="export_job_name=EXPORT&state_code=US_XX", ) self.assertEqual(HTTPStatus.OK, response.status_code) response = self.client.get( self.metric_view_data_export_url, headers=self.headers, query_string="export_job_name=export&state_code=us_xx", ) self.assertEqual(HTTPStatus.OK, response.status_code) @mock.patch( "recidiviz.metrics.export.view_export_manager.export_view_data_to_cloud_storage" ) def test_metric_view_data_export_state_agnostic( self, mock_export_view_data_to_cloud_storage: Mock) -> None: with self.app.test_request_context(): mock_export_view_data_to_cloud_storage.return_value = None response = self.client.get( self.metric_view_data_export_url, headers=self.headers, query_string="export_job_name=MOCK_EXPORT_NAME", ) self.assertEqual(HTTPStatus.OK, response.status_code) # case insensitive response = self.client.get( self.metric_view_data_export_url, headers=self.headers, query_string="export_job_name=mock_export_name", ) self.assertEqual(HTTPStatus.OK, response.status_code) @mock.patch( "recidiviz.metrics.export.view_export_manager.export_view_data_to_cloud_storage" ) def test_metric_view_data_export_missing_required_state_code( self, mock_export_view_data_to_cloud_storage: Mock) -> None: with self.app.test_request_context(): mock_export_view_data_to_cloud_storage.return_value = None response = self.client.get( self.metric_view_data_export_url, headers=self.headers, query_string="export_job_name=EXPORT", ) self.assertEqual(HTTPStatus.BAD_REQUEST, response.status_code) self.assertEqual( b"Missing required state_code parameter for export_job_name EXPORT", response.data, ) # case insensitive response = self.client.get( self.metric_view_data_export_url, headers=self.headers, query_string="export_job_name=export", ) self.assertEqual(HTTPStatus.BAD_REQUEST, response.status_code) self.assertEqual( b"Missing required state_code parameter for export_job_name EXPORT", response.data, ) @mock.patch( "recidiviz.metrics.export.view_export_manager.export_view_data_to_cloud_storage" ) def test_metric_view_data_export_missing_export_job_name( self, mock_export_view_data_to_cloud_storage: Mock) -> None: with self.app.test_request_context(): mock_export_view_data_to_cloud_storage.return_value = None response = self.client.get( self.metric_view_data_export_url, headers=self.headers, query_string="state_code=US_XX", ) self.assertEqual(HTTPStatus.BAD_REQUEST, response.status_code) self.assertEqual(b"Missing required export_job_name URL parameter", response.data) @mock.patch( "recidiviz.metrics.export.view_export_cloud_task_manager.ViewExportCloudTaskManager.create_metric_view_data_export_task" ) def test_create_metric_view_data_export_tasks_state_code_filter( self, mock_create_metric_view_data_export_task: Mock) -> None: with self.app.test_request_context(): mock_create_metric_view_data_export_task.return_value = None response = self.client.get( self.create_metric_view_data_export_tasks_url, headers=self.headers, query_string="export_job_filter=US_XX", ) self.assertEqual(HTTPStatus.OK, response.status_code) mock_create_metric_view_data_export_task.assert_has_calls( [ mock.call(export_job_name="EXPORT", state_code="US_XX"), mock.call(export_job_name="OTHER_EXPORT", state_code="US_XX"), ], any_order=True, ) response = self.client.get( self.create_metric_view_data_export_tasks_url, headers=self.headers, query_string="export_job_filter=us_xx", ) self.assertEqual(HTTPStatus.OK, response.status_code) mock_create_metric_view_data_export_task.assert_has_calls( [ mock.call(export_job_name="EXPORT", state_code="US_XX"), mock.call(export_job_name="OTHER_EXPORT", state_code="US_XX"), ], any_order=True, ) @mock.patch( "recidiviz.metrics.export.view_export_cloud_task_manager.ViewExportCloudTaskManager.create_metric_view_data_export_task" ) def test_create_metric_view_data_export_tasks_export_name_filter_state_agnostic( self, mock_create_metric_view_data_export_task: Mock) -> None: with self.app.test_request_context(): mock_create_metric_view_data_export_task.return_value = None response = self.client.get( self.create_metric_view_data_export_tasks_url, headers=self.headers, query_string="export_job_filter=MOCK_EXPORT_NAME", ) self.assertEqual(HTTPStatus.OK, response.status_code) mock_create_metric_view_data_export_task.assert_has_calls( [ mock.call(export_job_name="MOCK_EXPORT_NAME", state_code=None), ], any_order=True, ) # case insensitive response = self.client.get( self.create_metric_view_data_export_tasks_url, headers=self.headers, query_string="export_job_filter=mock_export_name", ) self.assertEqual(HTTPStatus.OK, response.status_code) mock_create_metric_view_data_export_task.assert_has_calls( [ mock.call(export_job_name="MOCK_EXPORT_NAME", state_code=None), ], any_order=True, ) @mock.patch( "recidiviz.metrics.export.view_export_cloud_task_manager.ViewExportCloudTaskManager.create_metric_view_data_export_task" ) def test_create_metric_view_data_export_tasks_export_name_filter( self, mock_create_metric_view_data_export_task: Mock) -> None: with self.app.test_request_context(): mock_create_metric_view_data_export_task.return_value = None response = self.client.get( self.create_metric_view_data_export_tasks_url, headers=self.headers, query_string="export_job_filter=EXPORT", ) self.assertEqual(HTTPStatus.OK, response.status_code) mock_create_metric_view_data_export_task.assert_has_calls( [ mock.call(export_job_name="EXPORT", state_code="US_XX"), mock.call(export_job_name="EXPORT", state_code="US_WW"), ], any_order=True, ) # case insensitive response = self.client.get( self.create_metric_view_data_export_tasks_url, headers=self.headers, query_string="export_job_filter=export", ) self.assertEqual(HTTPStatus.OK, response.status_code) mock_create_metric_view_data_export_task.assert_has_calls( [ mock.call(export_job_name="EXPORT", state_code="US_XX"), mock.call(export_job_name="EXPORT", state_code="US_WW"), ], any_order=True, )
class TestExportViewCollectionConfig(unittest.TestCase): """Tests the functionality of the ExportViewCollectionConfig class.""" def setUp(self): self.mock_project_id = "fake-recidiviz-project" self.mock_dataset_id = "base_dataset" self.mock_dataset = bigquery.dataset.DatasetReference( self.mock_project_id, self.mock_dataset_id ) self.metadata_patcher = mock.patch("recidiviz.utils.metadata.project_id") self.mock_project_id_fn = self.metadata_patcher.start() self.mock_project_id_fn.return_value = self.mock_project_id self.mock_big_query_view_namespace = BigQueryViewNamespace.STATE self.mock_view_builder = MetricBigQueryViewBuilder( dataset_id=self.mock_dataset.dataset_id, view_id="test_view", view_query_template="SELECT NULL LIMIT 0", dimensions=[], ) self.views_for_dataset = [self.mock_view_builder] def tearDown(self): self.metadata_patcher.stop() def test_matches_filter(self): """Tests matches_filter function to ensure that state codes and export names correctly match""" state_dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter="US_XX", export_name="EXPORT", bq_view_namespace=self.mock_big_query_view_namespace, ) self.assertTrue(state_dataset_export_config.matches_filter("US_XX")) dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter=None, export_name="VALID_EXPORT_NAME", bq_view_namespace=self.mock_big_query_view_namespace, ) self.assertTrue(dataset_export_config.matches_filter("VALID_EXPORT_NAME")) self.assertFalse(dataset_export_config.matches_filter("INVALID_EXPORT_NAME")) def test_matches_filter_case_insensitive(self): """Tests matches_filter function with different cases to ensure state codes and export names correctly match""" state_dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter="US_XX", export_name="OTHER_EXPORT", bq_view_namespace=self.mock_big_query_view_namespace, ) self.assertTrue(state_dataset_export_config.matches_filter("US_xx")) dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter=None, export_name="VALID_EXPORT_NAME", bq_view_namespace=self.mock_big_query_view_namespace, ) self.assertTrue(dataset_export_config.matches_filter("valid_export_name")) def test_metric_export_state_agnostic(self): """Tests the export_configs_for_views_to_export function on the ExportViewCollectionConfig class when the export is state-agnostic.""" state_agnostic_dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket-without-state-codes", state_code_filter=None, export_name="ALL_STATE_TEST_PRODUCT", bq_view_namespace=self.mock_big_query_view_namespace, ) view_configs_to_export = ( state_agnostic_dataset_export_config.export_configs_for_views_to_export( project_id=self.mock_project_id ) ) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportBigQueryViewConfig( view=expected_view, view_filter_clause=None, intermediate_table_name=f"{expected_view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( state_agnostic_dataset_export_config.output_directory_uri_template.format( project_id=self.mock_project_id, ) ), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ) ] self.assertEqual(expected_view_export_configs, view_configs_to_export) def test_metric_export_state_specific(self): """Tests the export_configs_for_views_to_export function on the ExportViewCollectionConfig class when the export is state-specific.""" specific_state_dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter="US_XX", export_name="TEST_REPORT", bq_view_namespace=self.mock_big_query_view_namespace, ) view_configs_to_export = ( specific_state_dataset_export_config.export_configs_for_views_to_export( project_id=self.mock_project_id ) ) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportBigQueryViewConfig( view=expected_view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{expected_view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( f"gs://{self.mock_project_id}-bucket/US_XX" ), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ) ] self.assertEqual(expected_view_export_configs, view_configs_to_export) def test_metric_export_lantern_dashboard(self): """Tests the export_configs_for_views_to_export function on the ExportViewCollectionConfig class when the export is state-agnostic.""" lantern_dashboard_dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket-without-state-codes", state_code_filter=None, export_name="TEST_EXPORT", bq_view_namespace=self.mock_big_query_view_namespace, ) view_configs_to_export = ( lantern_dashboard_dataset_export_config.export_configs_for_views_to_export( project_id=self.mock_project_id ) ) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportBigQueryViewConfig( view=expected_view, view_filter_clause=None, intermediate_table_name=f"{expected_view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( lantern_dashboard_dataset_export_config.output_directory_uri_template.format( project_id=self.mock_project_id, ) ), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ) ] self.assertEqual(expected_view_export_configs, view_configs_to_export) def test_metric_export_lantern_dashboard_with_state(self): """Tests the export_configs_for_views_to_export function on the ExportViewCollectionConfig class when the export is state-specific.""" lantern_dashboard_with_state_dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter="US_XX", export_name="TEST_EXPORT", bq_view_namespace=self.mock_big_query_view_namespace, ) view_configs_to_export = lantern_dashboard_with_state_dataset_export_config.export_configs_for_views_to_export( project_id=self.mock_project_id ) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportBigQueryViewConfig( view=expected_view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{expected_view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( f"gs://{self.mock_project_id}-bucket/US_XX" ), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ) ] self.assertEqual(expected_view_export_configs, view_configs_to_export)
class TestExportMetricDatasetConfig(unittest.TestCase): """Tests the functionality of the ExportMetricDatasetConfig class.""" def setUp(self): self.mock_project_id = 'fake-recidiviz-project' self.mock_dataset_id = 'base_dataset' self.mock_dataset = bigquery.dataset.DatasetReference( self.mock_project_id, self.mock_dataset_id) self.metadata_patcher = mock.patch( 'recidiviz.utils.metadata.project_id') self.mock_project_id_fn = self.metadata_patcher.start() self.mock_project_id_fn.return_value = self.mock_project_id self.mock_view_builder = MetricBigQueryViewBuilder( dataset_id=self.mock_dataset.dataset_id, view_id='test_view', view_query_template='SELECT NULL LIMIT 0', dimensions=[]) self.views_for_dataset = [self.mock_view_builder] def tearDown(self): self.metadata_patcher.stop() def test_matches_filter(self): """Tests matches_filter function to ensure that state codes and export names correctly match""" state_dataset_export_config = ExportMetricDatasetConfig( dataset_id='dataset_id', metric_view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter='US_XX', export_name=None) self.assertTrue(state_dataset_export_config.matches_filter('US_XX')) dataset_export_config = ExportMetricDatasetConfig( dataset_id='dataset_id', metric_view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter=None, export_name='VALID_EXPORT_NAME') self.assertTrue( dataset_export_config.matches_filter('VALID_EXPORT_NAME')) self.assertFalse( dataset_export_config.matches_filter('INVALID_EXPORT_NAME')) def test_matches_filter_case_insensitive(self): """Tests matches_filter function with different cases to ensure state codes and export names correctly match""" state_dataset_export_config = ExportMetricDatasetConfig( dataset_id='dataset_id', metric_view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter='US_XX', export_name=None) self.assertTrue(state_dataset_export_config.matches_filter('US_xx')) dataset_export_config = ExportMetricDatasetConfig( dataset_id='dataset_id', metric_view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter=None, export_name='VALID_EXPORT_NAME') self.assertTrue( dataset_export_config.matches_filter('valid_export_name')) def test_metric_export_state_agnostic(self): """Tests the export_configs_for_views_to_export function on the ExportMetricDatasetConfig class when the export is state-agnostic.""" state_agnostic_dataset_export_config = ExportMetricDatasetConfig( dataset_id='dataset_id', metric_view_builders_to_export=self.views_for_dataset, output_directory_uri_template= "gs://{project_id}-bucket-without-state-codes", state_code_filter=None, export_name=None) view_configs_to_export = state_agnostic_dataset_export_config.export_configs_for_views_to_export( project_id=self.mock_project_id) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportMetricBigQueryViewConfig( view=expected_view, view_filter_clause=None, intermediate_table_name=f"{expected_view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( state_agnostic_dataset_export_config. output_directory_uri_template.format( project_id=self.mock_project_id, ))) ] self.assertEqual(expected_view_export_configs, view_configs_to_export) def test_metric_export_state_specific(self): """Tests the export_configs_for_views_to_export function on the ExportMetricDatasetConfig class when the export is state-specific.""" specific_state_dataset_export_config = ExportMetricDatasetConfig( dataset_id='dataset_id', metric_view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter='US_XX', export_name=None) view_configs_to_export = specific_state_dataset_export_config.export_configs_for_views_to_export( project_id=self.mock_project_id) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportMetricBigQueryViewConfig( view=expected_view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{expected_view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( f"gs://{self.mock_project_id}-bucket/US_XX")) ] self.assertEqual(expected_view_export_configs, view_configs_to_export) def test_metric_export_lantern_dashboard(self): """Tests the export_configs_for_views_to_export function on the ExportMetricDatasetConfig class when the export is state-agnostic.""" lantern_dashboard_dataset_export_config = ExportMetricDatasetConfig( dataset_id='dataset_id', metric_view_builders_to_export=self.views_for_dataset, output_directory_uri_template= "gs://{project_id}-bucket-without-state-codes", state_code_filter=None, export_name="TEST_EXPORT") view_configs_to_export = lantern_dashboard_dataset_export_config.export_configs_for_views_to_export( project_id=self.mock_project_id) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportMetricBigQueryViewConfig( view=expected_view, view_filter_clause=None, intermediate_table_name=f"{expected_view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( lantern_dashboard_dataset_export_config. output_directory_uri_template.format( project_id=self.mock_project_id, ))) ] self.assertEqual(expected_view_export_configs, view_configs_to_export) def test_metric_export_lantern_dashboard_with_state(self): """Tests the export_configs_for_views_to_export function on the ExportMetricDatasetConfig class when the export is state-specific.""" lantern_dashboard_with_state_dataset_export_config = ExportMetricDatasetConfig( dataset_id='dataset_id', metric_view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", state_code_filter="US_XX", export_name="TEST_EXPORT") view_configs_to_export = lantern_dashboard_with_state_dataset_export_config.export_configs_for_views_to_export( project_id=self.mock_project_id) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportMetricBigQueryViewConfig( view=expected_view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{expected_view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( f"gs://{self.mock_project_id}-bucket/US_XX")) ] self.assertEqual(expected_view_export_configs, view_configs_to_export)
class TestExportViewCollectionConfig(unittest.TestCase): """Tests the functionality of the ExportViewCollectionConfig class.""" def setUp(self) -> None: self.mock_project_id = "fake-recidiviz-project" self.mock_dataset_id = "base_dataset" self.mock_dataset = bigquery.dataset.DatasetReference( self.mock_project_id, self.mock_dataset_id) self.metadata_patcher = mock.patch( "recidiviz.utils.metadata.project_id") self.mock_project_id_fn = self.metadata_patcher.start() self.mock_project_id_fn.return_value = self.mock_project_id self.mock_big_query_view_namespace = BigQueryViewNamespace.STATE self.mock_view_builder = MetricBigQueryViewBuilder( dataset_id=self.mock_dataset.dataset_id, view_id="test_view", description="test_view description", view_query_template="SELECT NULL LIMIT 0", dimensions=(), ) self.views_for_dataset = [self.mock_view_builder] def tearDown(self) -> None: self.metadata_patcher.stop() def test_unique_export_names(self) -> None: self.assertEqual( len(_VIEW_COLLECTION_EXPORT_CONFIGS), len(VIEW_COLLECTION_EXPORT_INDEX.keys()), ) def test_metric_export_state_agnostic(self) -> None: """Tests the export_configs_for_views_to_export function on the ExportViewCollectionConfig class when the export is state-agnostic.""" state_agnostic_dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template= "gs://{project_id}-bucket-without-state-codes", export_name="ALL_STATE_TEST_PRODUCT", bq_view_namespace=self.mock_big_query_view_namespace, ) view_configs_to_export = (state_agnostic_dataset_export_config. export_configs_for_views_to_export( project_id=self.mock_project_id, )) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=expected_view, view_filter_clause=None, intermediate_table_name=f"{expected_view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( state_agnostic_dataset_export_config. output_directory_uri_template.format( project_id=self.mock_project_id, )), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ) ] self.assertEqual(expected_view_export_configs, view_configs_to_export) def test_metric_export_state_specific(self) -> None: """Tests the export_configs_for_views_to_export function on the ExportViewCollectionConfig class when the export is state-specific.""" specific_state_dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", export_name="STATE_SPECIFIC_PRODUCT_EXPORT", bq_view_namespace=self.mock_big_query_view_namespace, ) mock_export_job_filter = "US_XX" view_configs_to_export = (specific_state_dataset_export_config. export_configs_for_views_to_export( project_id=self.mock_project_id, state_code_filter=mock_export_job_filter, )) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=expected_view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{expected_view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( f"gs://{self.mock_project_id}-bucket/US_XX"), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ) ] self.assertEqual(expected_view_export_configs, view_configs_to_export) def test_metric_export_lantern_dashboard(self) -> None: """Tests the export_configs_for_views_to_export function on the ExportViewCollectionConfig class when the export is state-agnostic.""" lantern_dashboard_dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template= "gs://{project_id}-bucket-without-state-codes", export_name="TEST_EXPORT", bq_view_namespace=self.mock_big_query_view_namespace, ) view_configs_to_export = (lantern_dashboard_dataset_export_config. export_configs_for_views_to_export( project_id=self.mock_project_id, )) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=expected_view, view_filter_clause=None, intermediate_table_name=f"{expected_view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( lantern_dashboard_dataset_export_config. output_directory_uri_template.format( project_id=self.mock_project_id, )), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ) ] self.assertEqual(expected_view_export_configs, view_configs_to_export) def test_metric_export_lantern_dashboard_with_state(self) -> None: """Tests the export_configs_for_views_to_export function on the ExportViewCollectionConfig class when the export is state-specific.""" lantern_dashboard_with_state_dataset_export_config = ExportViewCollectionConfig( view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-bucket", export_name="TEST_EXPORT", bq_view_namespace=self.mock_big_query_view_namespace, ) mock_state_code = "US_XX" view_configs_to_export = lantern_dashboard_with_state_dataset_export_config.export_configs_for_views_to_export( project_id=self.mock_project_id, state_code_filter=mock_state_code) expected_view = self.mock_view_builder.build() expected_view_export_configs = [ ExportBigQueryViewConfig( bq_view_namespace=self.mock_big_query_view_namespace, view=expected_view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{expected_view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( f"gs://{self.mock_project_id}-bucket/US_XX"), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ) ] self.assertEqual(expected_view_export_configs, view_configs_to_export)
class MetricViewExportManagerTest(unittest.TestCase): """Tests for metric_view_export_manager.py.""" def setUp(self): self.mock_project_id = 'fake-recidiviz-project' self.mock_dataset_id = 'base_dataset' self.mock_dataset = bigquery.dataset.DatasetReference( self.mock_project_id, self.mock_dataset_id) self.metadata_patcher = mock.patch('recidiviz.utils.metadata.project_id') self.mock_project_id_fn = self.metadata_patcher.start() self.mock_project_id_fn.return_value = self.mock_project_id self.client_patcher = mock.patch( 'recidiviz.metrics.export.metric_view_export_manager.BigQueryClientImpl') self.mock_client = self.client_patcher.start().return_value self.mock_client.dataset_ref_for_id.return_value = self.mock_dataset self.mock_view_builder = MetricBigQueryViewBuilder(dataset_id=self.mock_dataset.dataset_id, view_id='test_view', view_query_template='SELECT NULL LIMIT 0', dimensions=[]) self.views_for_dataset = [self.mock_view_builder] self.output_uri_template_for_dataset = { "dataset_id": "gs://{project_id}-dataset-location/subdirectory", } self.views_to_update = {self.mock_dataset_id: self.views_for_dataset} self.metric_dataset_export_configs = [ ExportMetricDatasetConfig( dataset_id=self.mock_dataset_id, metric_view_builders_to_export=self.views_for_dataset, output_directory_uri_template="gs://{project_id}-dataset-location/subdirectory", state_code_filter=mock_state_code, export_name=None ) ] view_config_values = { 'OUTPUT_DIRECTORY_URI_TEMPLATE_FOR_DATASET_EXPORT': self.output_uri_template_for_dataset, 'VIEW_BUILDERS_FOR_VIEWS_TO_UPDATE': self.views_to_update, 'METRIC_DATASET_EXPORT_CONFIGS': self.metric_dataset_export_configs } self.view_export_config_patcher = mock.patch( 'recidiviz.metrics.export.metric_view_export_manager.view_config', **view_config_values) self.mock_export_config = self.view_export_config_patcher.start() def tearDown(self): self.client_patcher.stop() self.view_export_config_patcher.stop() self.metadata_patcher.stop() @mock.patch('recidiviz.big_query.view_update_manager.create_dataset_and_update_views_for_view_builders') @mock.patch('recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter') def test_export_dashboard_data_to_cloud_storage(self, mock_view_exporter, mock_view_update_manager): """Tests the table is created from the view and then extracted.""" metric_view_export_manager.export_view_data_to_cloud_storage(mock_state_code, mock_view_exporter) view = self.mock_view_builder.build() view_export_configs = [ExportMetricBigQueryViewConfig( view=view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory/{state_code}".format( project_id=self.mock_project_id, state_code='US_XX', ) ) )] mock_view_update_manager.assert_called() mock_view_exporter.export_and_validate.assert_called_with(view_export_configs) @mock.patch('recidiviz.big_query.view_update_manager.create_dataset_and_update_views_for_view_builders') @mock.patch('recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter') def test_raise_exception_no_export_matched(self, mock_view_exporter, mock_view_update_manager): # pylint: disable=unused-argument """Tests the table is created from the view and then extracted.""" with self.assertRaises(ValueError) as e: metric_view_export_manager.export_view_data_to_cloud_storage('US_YY', mock_view_exporter) self.assertEqual(str(e.exception), 'Export filter did not match any export configs:', ' US_YY') @mock.patch('recidiviz.metrics.export.metric_view_export_manager.view_config') @mock.patch('recidiviz.big_query.view_update_manager.create_dataset_and_update_views_for_view_builders') @mock.patch('recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter') def test_export_dashboard_data_to_cloud_storage_state_agnostic(self, mock_view_exporter, mock_view_update_manager, mock_view_config): """Tests the table is created from the view and then extracted, where the export is not state-specific.""" state_agnostic_dataset_export_configs = [ ExportMetricDatasetConfig( dataset_id='dataset_id', metric_view_builders_to_export=[self.mock_view_builder], output_directory_uri_template="gs://{project_id}-bucket-without-state-codes", state_code_filter=None, export_name=None ), ] mock_view_config.VIEW_BUILDERS_FOR_VIEWS_TO_UPDATE = self.views_to_update mock_view_config.METRIC_DATASET_EXPORT_CONFIGS = state_agnostic_dataset_export_configs metric_view_export_manager.export_view_data_to_cloud_storage(export_job_filter=None, view_exporter=mock_view_exporter) view = self.mock_view_builder.build() view_export_configs = [ExportMetricBigQueryViewConfig( view=view, view_filter_clause=None, intermediate_table_name=f"{view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-bucket-without-state-codes".format( project_id=self.mock_project_id, ) ) )] mock_view_update_manager.assert_called() mock_view_exporter.export_and_validate.assert_called_with(view_export_configs) @mock.patch('recidiviz.big_query.view_update_manager.create_dataset_and_update_views_for_view_builders') @mock.patch('recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter') def test_export_dashboard_data_to_cloud_storage_value_error(self, mock_view_exporter, mock_view_update_manager): """Tests the table is created from the view and then extracted.""" mock_view_exporter.export_and_validate.side_effect = ValueError with self.assertRaises(ValueError): metric_view_export_manager.export_view_data_to_cloud_storage(mock_state_code, mock_view_exporter) view = self.mock_view_builder.build() view_export_configs = [ExportMetricBigQueryViewConfig( view=view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory/{state_code}".format( project_id=self.mock_project_id, state_code='US_XX', ) ) )] mock_view_update_manager.assert_called() mock_view_exporter.export_and_validate.assert_called_with(view_export_configs) @mock.patch('recidiviz.big_query.view_update_manager.create_dataset_and_update_views_for_view_builders') @mock.patch('recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter') def test_export_dashboard_data_to_cloud_storage_validation_error(self, mock_view_exporter, mock_view_update_manager): """Tests the table is created from the view and then extracted.""" mock_view_exporter.export_and_validate.side_effect = ViewExportValidationError # Should not throw metric_view_export_manager.export_view_data_to_cloud_storage(mock_state_code, mock_view_exporter) view = self.mock_view_builder.build() view_export_configs = [ExportMetricBigQueryViewConfig( view=view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory/{state_code}".format( project_id=self.mock_project_id, state_code='US_XX', ) ) )] mock_view_update_manager.assert_called() mock_view_exporter.export_and_validate.assert_called_with(view_export_configs)
class ViewCollectionExportManagerTest(unittest.TestCase): """Tests for view_export_manager.py.""" def setUp(self) -> None: self.mock_state_code = "US_XX" self.mock_project_id = "fake-recidiviz-project" self.mock_dataset_id = "base_dataset" self.mock_dataset = bigquery.dataset.DatasetReference( self.mock_project_id, self.mock_dataset_id) self.metadata_patcher = mock.patch( "recidiviz.utils.metadata.project_id") self.mock_project_id_fn = self.metadata_patcher.start() self.mock_project_id_fn.return_value = self.mock_project_id self.client_patcher = mock.patch( "recidiviz.metrics.export.view_export_manager.BigQueryClientImpl") self.mock_client = self.client_patcher.start().return_value self.mock_client.dataset_ref_for_id.return_value = self.mock_dataset self.mock_view_builder = SimpleBigQueryViewBuilder( dataset_id=self.mock_dataset.dataset_id, view_id="test_view", view_query_template="SELECT NULL LIMIT 0", ) self.mock_metric_view_builder = MetricBigQueryViewBuilder( dataset_id=self.mock_dataset.dataset_id, view_id="test_view", view_query_template="SELECT NULL LIMIT 0", dimensions=tuple(), ) self.view_buidlers_for_dataset = [ self.mock_view_builder, self.mock_metric_view_builder, ] self.output_uri_template_for_dataset = { "dataset_id": "gs://{project_id}-dataset-location/subdirectory", } self.views_to_update = { self.mock_dataset_id: self.view_buidlers_for_dataset } self.mock_export_name = "MOCK_EXPORT_NAME" self.mock_big_query_view_namespace = BigQueryViewNamespace.STATE self.metric_dataset_export_configs = [ ExportViewCollectionConfig( view_builders_to_export=self.view_buidlers_for_dataset, output_directory_uri_template= "gs://{project_id}-dataset-location/subdirectory", state_code_filter=self.mock_state_code, export_name=self.mock_export_name, bq_view_namespace=self.mock_big_query_view_namespace, ) ] export_config_values = { "OUTPUT_DIRECTORY_URI_TEMPLATE_FOR_DATASET_EXPORT": self.output_uri_template_for_dataset, "VIEW_COLLECTION_EXPORT_CONFIGS": self.metric_dataset_export_configs, } self.export_config_patcher = mock.patch( # type: ignore[call-overload] "recidiviz.metrics.export.view_export_manager.export_config", **export_config_values, ) self.mock_export_config = self.export_config_patcher.start() self.gcs_factory_patcher = mock.patch( "recidiviz.metrics.export.view_export_manager.GcsfsFactory.build") self.gcs_factory_patcher.start().return_value = FakeGCSFileSystem() def tearDown(self) -> None: self.client_patcher.stop() self.export_config_patcher.stop() self.metadata_patcher.stop() self.gcs_factory_patcher.stop() @mock.patch( "recidiviz.big_query.view_update_manager.rematerialize_views_for_namespace" ) @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage( self, mock_view_exporter, mock_view_update_manager_rematerialize) -> None: """Tests the table is created from the view and then extracted.""" view_export_manager.export_view_data_to_cloud_storage( self.mock_state_code, mock_view_exporter) view = self.mock_view_builder.build() metric_view = self.mock_metric_view_builder.build() view_export_configs = [ ExportBigQueryViewConfig( view=view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory/{state_code}" .format( project_id=self.mock_project_id, state_code="US_XX", )), export_output_formats=[ExportOutputFormatType.JSON], ), ExportBigQueryViewConfig( view=metric_view, view_filter_clause=" WHERE state_code = 'US_XX'", intermediate_table_name=f"{view.view_id}_table_US_XX", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-dataset-location/subdirectory/{state_code}" .format( project_id=self.mock_project_id, state_code="US_XX", )), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ), ] mock_view_update_manager_rematerialize.assert_called() mock_view_exporter.export_and_validate.assert_has_calls( [ mock.call([]), # CSV export mock.call([ view_export_configs[1].pointed_to_staging_subdirectory() ]), # JSON export mock.call([ conf.pointed_to_staging_subdirectory() for conf in view_export_configs ]), # METRIC export ], any_order=True, ) @mock.patch( "recidiviz.big_query.view_update_manager.create_dataset_and_deploy_views_for_view_builders" ) @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_raise_exception_no_export_matched( self, mock_view_exporter, mock_view_update_manager_rematerialize) -> None: # pylint: disable=unused-argument """Tests the table is created from the view and then extracted.""" self.mock_export_config.NAMESPACE_TO_UPDATE_FOR_EXPORT_FILTER = { "US_YY": "NAMESPACE" } with self.assertRaises(ValueError) as e: view_export_manager.export_view_data_to_cloud_storage( "US_YY", mock_view_exporter) self.assertEqual( str(e.exception), "Export filter did not match any export configs:", " US_YY", ) @mock.patch( "recidiviz.big_query.view_update_manager.rematerialize_views_for_namespace" ) @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage_state_agnostic( self, mock_view_exporter, mock_view_update_manager_rematerialize) -> None: """Tests the table is created from the view and then extracted, where the export is not state-specific.""" state_agnostic_dataset_export_configs = [ ExportViewCollectionConfig( view_builders_to_export=self.view_buidlers_for_dataset, output_directory_uri_template= "gs://{project_id}-bucket-without-state-codes", state_code_filter=None, export_name=self.mock_export_name, bq_view_namespace=self.mock_big_query_view_namespace, ), ] self.mock_export_config.VIEW_COLLECTION_EXPORT_CONFIGS = ( state_agnostic_dataset_export_configs) view_export_manager.export_view_data_to_cloud_storage( export_job_filter=self.mock_export_name, override_view_exporter=mock_view_exporter, ) view = self.mock_view_builder.build() metric_view = self.mock_metric_view_builder.build() view_export_configs = [ ExportBigQueryViewConfig( view=view, view_filter_clause=None, intermediate_table_name=f"{view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-bucket-without-state-codes".format( project_id=self.mock_project_id, )), export_output_formats=[ExportOutputFormatType.JSON], ), ExportBigQueryViewConfig( view=metric_view, view_filter_clause=None, intermediate_table_name=f"{view.view_id}_table", output_directory=GcsfsDirectoryPath.from_absolute_path( "gs://{project_id}-bucket-without-state-codes".format( project_id=self.mock_project_id, )), export_output_formats=[ ExportOutputFormatType.JSON, ExportOutputFormatType.METRIC, ], ), ] mock_view_update_manager_rematerialize.assert_called() mock_view_exporter.export_and_validate.assert_has_calls( [ mock.call([]), # CSV export mock.call([ view_export_configs[1].pointed_to_staging_subdirectory() ]), # JSON export mock.call([ conf.pointed_to_staging_subdirectory() for conf in view_export_configs ]), # METRIC export ], any_order=True, ) @mock.patch( "recidiviz.big_query.view_update_manager.rematerialize_views_for_namespace" ) @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage_value_error( self, mock_view_exporter, mock_view_update_manager_rematerialize) -> None: """Tests the table is created from the view and then extracted.""" mock_view_exporter.export_and_validate.side_effect = ValueError with self.assertRaises(ValueError): view_export_manager.export_view_data_to_cloud_storage( self.mock_state_code, mock_view_exporter) # Just the metric export is attempted and then the raise stops subsequent checks from happening mock_view_update_manager_rematerialize.assert_called_once() @mock.patch( "recidiviz.big_query.view_update_manager.rematerialize_views_for_namespace" ) @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage_validation_error( self, mock_view_exporter, mock_view_update_manager_rematerialize) -> None: """Tests the table is created from the view and then extracted.""" mock_view_exporter.export_and_validate.side_effect = ViewExportValidationError # Should not throw view_export_manager.export_view_data_to_cloud_storage( self.mock_state_code, mock_view_exporter) # Just the metric export is attempted and then the raise stops subsequent checks from happening mock_view_update_manager_rematerialize.assert_called_once() @mock.patch( "recidiviz.metrics.export.view_export_manager.view_update_manager.VIEW_BUILDERS_BY_NAMESPACE" ) @mock.patch( "recidiviz.big_query.view_update_manager.rematerialize_views_for_namespace" ) @mock.patch( "recidiviz.big_query.view_update_manager.create_dataset_and_deploy_views_for_view_builders" ) @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage_update_all_views( self, mock_view_exporter, mock_view_update_manager_deploy, mock_view_update_manager_rematerialize, mock_view_builders_by_namespace, ) -> None: """Tests that all views in the namespace are updated before the export when the export name is in export_config.NAMESPACES_REQUIRING_FULL_UPDATE.""" self.mock_export_config.NAMESPACES_REQUIRING_FULL_UPDATE = [ self.mock_big_query_view_namespace ] mock_view_builders_by_namespace.return_value = { self.mock_big_query_view_namespace: self.view_buidlers_for_dataset } view_export_manager.export_view_data_to_cloud_storage( self.mock_state_code, mock_view_exporter) mock_view_update_manager_deploy.assert_called_with( self.mock_big_query_view_namespace, mock_view_builders_by_namespace[ self.mock_big_query_view_namespace], ) mock_view_update_manager_rematerialize.assert_called_once() @mock.patch( "recidiviz.metrics.export.view_export_manager.view_update_manager.VIEW_BUILDERS_BY_NAMESPACE" ) @mock.patch( "recidiviz.big_query.view_update_manager.rematerialize_views_for_namespace" ) @mock.patch( "recidiviz.big_query.export.big_query_view_exporter.BigQueryViewExporter" ) def test_export_dashboard_data_to_cloud_storage_update_materialized_views_only( self, mock_view_exporter, mock_view_update_manager_rematerialize, mock_view_builders_by_namespace, ) -> None: """Tests that only materialized views in the namespace are updated before the export when the export name is not in export_config.NAMESPACES_REQUIRING_FULL_UPDATE.""" self.mock_export_config.NAMESPACES_REQUIRING_FULL_UPDATE = [ "OTHER_NAMESPACE" ] mock_view_builders_by_namespace.return_value = { self.mock_big_query_view_namespace: self.view_buidlers_for_dataset } view_export_manager.export_view_data_to_cloud_storage( self.mock_state_code, mock_view_exporter) mock_view_update_manager_rematerialize.assert_called_with( bq_view_namespace=self.mock_big_query_view_namespace, candidate_view_builders=mock_view_builders_by_namespace[ self.mock_big_query_view_namespace], )