From 202b73eb59bde5699f411bfb4c34818f9fe8a7b6 Mon Sep 17 00:00:00 2001 From: shubhamjagtap639 Date: Tue, 7 Nov 2023 17:15:52 +0530 Subject: [PATCH] Ingestion checkpoint provider test case modified --- .../provider/test_provider.py | 26 ++++++------------- 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/metadata-ingestion/tests/unit/stateful_ingestion/provider/test_provider.py b/metadata-ingestion/tests/unit/stateful_ingestion/provider/test_provider.py index 0006b28f627ba..4387e5a17790f 100644 --- a/metadata-ingestion/tests/unit/stateful_ingestion/provider/test_provider.py +++ b/metadata-ingestion/tests/unit/stateful_ingestion/provider/test_provider.py @@ -26,13 +26,10 @@ from datahub.ingestion.source.state_provider.file_ingestion_checkpointing_provider import ( FileIngestionCheckpointingProvider, ) -from datahub.ingestion.source.state_provider.state_provider_registry import ( - ingestion_checkpoint_provider_registry, -) from tests.test_helpers.type_helpers import assert_not_null -class TestDatahubIngestionCheckpointProvider(unittest.TestCase): +class TestIngestionCheckpointProviders(unittest.TestCase): # Static members for the tests pipeline_name: str = "test_pipeline" job_names: List[JobId] = [JobId("job1"), JobId("job2")] @@ -69,20 +66,13 @@ def _create_providers(self) -> None: run_id=self.run_id, pipeline_name=self.pipeline_name ) ctx.graph = self.mock_graph - self.providers: List[IngestionCheckpointingProviderBase] = [] - for provider_name in ingestion_checkpoint_provider_registry.mapping.keys(): - provider_cls = ingestion_checkpoint_provider_registry.get(provider_name) - if provider_cls == DatahubIngestionCheckpointingProvider: - self.providers.append( - DatahubIngestionCheckpointingProvider.create({}, ctx) - ) - elif provider_cls == FileIngestionCheckpointingProvider: - self.providers.append( - FileIngestionCheckpointingProvider.create( - {"filename": f"{tempfile.mkdtemp()}/checkpoint_mces.json"}, - ctx, - ) - ) + self.providers: List[IngestionCheckpointingProviderBase] = [ + DatahubIngestionCheckpointingProvider.create({}, ctx), + FileIngestionCheckpointingProvider.create( + {"filename": f"{tempfile.mkdtemp()}/checkpoint_mces.json"}, + ctx, + ), + ] def monkey_patch_emit_mcp( self, graph_ref: MagicMock, mcpw: MetadataChangeProposalWrapper