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)
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_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 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 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,
)
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()