def test_create_raw_data_latest_view_update_tasks( self, mock_cloud_task_manager_fn: mock.MagicMock, mock_environment: mock.MagicMock, ) -> None: with local_project_id_override("recidiviz-staging"): mock_environment.return_value = "staging" mock_cloud_task_manager = create_autospec( DirectIngestRawUpdateCloudTaskManager) mock_cloud_task_manager_fn.return_value = mock_cloud_task_manager headers = {"X-Appengine-Cron": "test-cron"} response = self.client.post( "/create_raw_data_latest_view_update_tasks", query_string={}, headers=headers, ) self.assertEqual(200, response.status_code) expected_calls = [ mock.call(region_code) for region_code in get_existing_region_dir_names() ] mock_cloud_task_manager.create_raw_data_latest_view_update_task.assert_has_calls( expected_calls)
def test_create_raw_data_latest_view_update_tasks( self, mock_cloud_task_manager_fn, mock_environment): with local_project_id_override('recidiviz-staging'): mock_environment.return_value = 'staging' mock_cloud_task_manager = create_autospec( DirectIngestRawUpdateCloudTaskManager) mock_cloud_task_manager_fn.return_value = mock_cloud_task_manager headers = {'X-Appengine-Cron': 'test-cron'} response = self.client.post( '/create_raw_data_latest_view_update_tasks', query_string={}, headers=headers) self.assertEqual(200, response.status_code) no_raw_data_import_regions = { 'us_tx_brazos', 'us_ma_middlesex', 'us_nm_bernalillo' } expected_calls = [ mock.call(region_code) for region_code in get_existing_region_dir_names() if region_code not in no_raw_data_import_regions ] mock_cloud_task_manager.create_raw_data_latest_view_update_task.assert_has_calls( expected_calls)
def test_collect_ingest_views(): with local_project_id_override('project'): for region_code in get_existing_region_dir_names(): region = get_region(region_code, is_direct_ingest=True) controller_class = region.get_ingestor_class() if not issubclass(controller_class, GcsfsDirectIngestController): continue _ = DirectIngestPreProcessedIngestViewCollector( region, controller_class.get_file_tag_rank_list()).collect_views()
def test_direct_ingest_instance_status_contains_data_for_all_states( self) -> None: '''Enforces that after all migrations the set of direct ingest instance statuses matches the list of known states. If this test fails, you will likely have to add a new migration because a new state was recently created. To do so, first run: ``` python -m recidiviz.tools.migrations.autogenerate_migration \ --database OPERATIONS \ --message add_us_xx ``` This will generate a blank migration. You should then modify the migration, changing the `upgrade` method to look like: ``` def upgrade() -> None: op.execute(""" INSERT INTO direct_ingest_instance_status (region_code, instance, is_paused) VALUES ('US_XX', 'PRIMARY', TRUE), ('US_XX', 'SECONDARY', TRUE); """) ``` Afterwards, this test should ideally pass. ''' with runner(self.default_config(), self.engine) as r: r.migrate_up_to("head") engine = create_engine( local_postgres_helpers.postgres_db_url_from_env_vars()) conn = engine.connect() rows = conn.execute( "SELECT region_code, instance FROM direct_ingest_instance_status;" ) instance_to_state_codes = defaultdict(set) for row in rows: instance_to_state_codes[DirectIngestInstance(row[1])].add( row[0]) required_states = { name.upper() for name in get_existing_region_dir_names() } for instance in DirectIngestInstance: self.assertEqual(required_states, instance_to_state_codes[instance])
def test_get_raw_data_configs(self) -> None: raw_data_configs = get_raw_data_configs("us_id") self.assertIsNotNone(raw_data_configs) mittimus = next(raw_data_config for raw_data_config in raw_data_configs if raw_data_config.file_tag == "mittimus") self.assertIn("mitt_srl", mittimus.columns) self.assertIn("mitt_srl", mittimus.primary_keys) self.assertEqual("|", mittimus.separator) self.assertEqual("ISO-8859-1", mittimus.encoding) self.assertEqual(csv.QUOTE_MINIMAL, mittimus.quoting) for region_code in get_existing_region_dir_names(): self.assertIsNotNone(get_raw_data_configs(region_code))
def test_raw_files_yaml_parses_all_regions(self): for region_code in get_existing_region_dir_names(): region = get_region(region_code, is_direct_ingest=True) raw_file_manager = DirectIngestRegionRawFileConfig( region_code=region.region_code) if region.raw_data_bq_imports_enabled_env is not None: self.assertTrue(raw_file_manager.raw_file_configs) config_file_tags = set() for config in raw_file_manager.raw_file_configs.values(): self.assertTrue( config.file_tag not in config_file_tags, f"Multiple raw file configs defined with the same file_tag [{config.file_tag}]" ) config_file_tags.add(config.file_tag)
def test_build_gcsfs_ingest_controller_all_regions(self) -> None: for region_code in get_existing_region_dir_names(): region = get_region(region_code, is_direct_ingest=True) for ingest_instance in DirectIngestInstance: ingest_bucket_path = gcsfs_direct_ingest_bucket_for_region( region_code=region_code, system_level=SystemLevel.for_region(region), ingest_instance=ingest_instance, ) controller = DirectIngestControllerFactory.build( ingest_bucket_path=ingest_bucket_path, allow_unlaunched=False) self.assertIsNotNone(controller) self.assertIsInstance(controller, BaseDirectIngestController) self.assertEqual(ingest_bucket_path, controller.ingest_bucket_path)
def collect_view_builders( self) -> List[DirectIngestRawDataTableLatestViewBuilder]: builder_list = [] for region_code in get_existing_region_dir_names(): region_raw_file_config = DirectIngestRegionRawFileConfig( region_code) raw_file_configs = region_raw_file_config.raw_file_configs builder_list.extend([ DirectIngestRawDataTableLatestViewBuilder( region_code=region_code, raw_file_config=config, should_build_predicate=BigQueryTableChecker( f"{region_raw_file_config.region_code.lower()}_raw_data", config.file_tag, ).get_table_exists_predicate(), ) for config in raw_file_configs.values() if not config.is_undocumented and config.primary_key_cols ]) return builder_list
def test_build_gcsfs_ingest_controller_all_regions_do_not_allow_launched( self, ) -> None: for region_code in get_existing_region_dir_names(): region = get_region(region_code, is_direct_ingest=True) for ingest_instance in DirectIngestInstance: ingest_bucket_path = gcsfs_direct_ingest_bucket_for_region( region_code=region_code, system_level=SystemLevel.for_region(region), ingest_instance=ingest_instance, ) controller = DirectIngestControllerFactory.build( ingest_bucket_path=ingest_bucket_path, allow_unlaunched=True) # Should still succeed for all controllers in the test environment self.assertIsNotNone(controller) self.assertIsInstance(controller, BaseDirectIngestController) self.assertEqual(ingest_bucket_path, controller.ingest_bucket_path)
def create_raw_data_latest_view_update_tasks() -> Tuple[str, HTTPStatus]: """Creates tasks for every direct ingest region with SQL preprocessing enabled to update the raw data table latest views. """ raw_update_ctm = DirectIngestRawUpdateCloudTaskManager( metadata.project_id()) for region_code in get_existing_region_dir_names(): with monitoring.push_region_tag(region_code): region = get_region(region_code, is_direct_ingest=True) if region.are_raw_data_bq_imports_enabled_in_env(): logging.info( 'Creating raw data latest view update task for region [%s]', region_code) raw_update_ctm.create_raw_data_latest_view_update_task( region_code) else: logging.info( 'Skipping raw data latest view update for region [%s] - raw data imports not enabled.', region_code) return '', HTTPStatus.OK
def create_raw_data_latest_view_update_tasks() -> Tuple[str, HTTPStatus]: """Creates tasks for every direct ingest region with SQL preprocessing enabled to update the raw data table latest views. """ raw_update_ctm = DirectIngestRawUpdateCloudTaskManager() for region_code in get_existing_region_dir_names(): with monitoring.push_region_tag(region_code, ingest_instance=None): region = _region_for_region_code(region_code) if region.is_ingest_launched_in_env(): logging.info( "Creating raw data latest view update task for region [%s]", region_code, ) raw_update_ctm.create_raw_data_latest_view_update_task( region_code) else: logging.info( "Skipping raw data latest view update for region [%s] - ingest not enabled.", region_code, ) return "", HTTPStatus.OK
def region_dir_paths(self) -> List[str]: return [ os.path.join(self.temp_dir, REGIONS_DIR_PATH, d) for d in get_existing_region_dir_names() ]
def region_dir_names(self) -> List[str]: return get_existing_region_dir_names()
def test_region_dirname_matches_pattern(self): for d in get_existing_region_dir_names(): self.assertIsNotNone( re.match(_REGION_REGEX, d), f'Region [{d}] does not match expected region pattern.')