Fixing tests

datahub-project · Oct 3, 2023 · c52ea07 · c52ea07
1 parent 461e8d8
commit c52ea07
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 19 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/api/source.py b/metadata-ingestion/src/datahub/ingestion/api/source.py
@@ -196,16 +196,15 @@ def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
                 self.ctx.pipeline_config.flags.generate_browse_path_v2_dry_run
             )
 
-        auto_lowercase_urns: Optional[MetadataWorkUnitProcessor] = None
-        if self.ctx.pipeline_config and self.ctx.pipeline_config.source.config:
-            auto_lowercase_urns = self._get_auto_lowercase_urn_processor(
-                enabled=self.ctx.pipeline_config.source.config.get(
-                    LOWER_CASE_URN_CONFIG_KEY, False
-                )
-            )
-
+        auto_lowercase_dataset_urns: Optional[MetadataWorkUnitProcessor] = None
+        if (
+            self.ctx.pipeline_config
+            and self.ctx.pipeline_config.source.config
+            and self.ctx.pipeline_config.source.config.get(LOWER_CASE_URN_CONFIG_KEY)
+        ):
+            auto_lowercase_dataset_urns = auto_lowercase_urns
         return [
-            auto_lowercase_urns,
+            auto_lowercase_dataset_urns,
             auto_status_aspect,
             auto_materialize_referenced_tags,
             browse_path_processor,
@@ -251,11 +250,6 @@ def get_report(self) -> SourceReport:
     def close(self) -> None:
         pass
 
-    def _get_auto_lowercase_urn_processor(
-        self, enabled: bool
-    ) -> MetadataWorkUnitProcessor:
-        return partial(auto_lowercase_urns, enabled=enabled)
-
     def _get_browse_path_processor(self, dry_run: bool) -> MetadataWorkUnitProcessor:
         config = self.get_config()
         platform = getattr(self, "platform", None) or getattr(config, "platform", None)

diff --git a/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py b/metadata-ingestion/src/datahub/ingestion/api/source_helpers.py
@@ -174,13 +174,10 @@ def auto_materialize_referenced_tags(
 
 
 def auto_lowercase_urns(
-    stream: Iterable[MetadataWorkUnit], enabled: bool = False
+    stream: Iterable[MetadataWorkUnit],
 ) -> Iterable[MetadataWorkUnit]:
     """Lowercase all dataset urns"""
 
-    if not enabled:
-        return stream
-
     for wu in stream:
         try:
             old_urn = wu.get_urn()

diff --git a/metadata-ingestion/tests/unit/test_source_helpers.py b/metadata-ingestion/tests/unit/test_source_helpers.py
@@ -342,7 +342,7 @@ def test_auto_lowercase_aspects():
             )
         ),
     ]
-    assert list(auto_lowercase_urns(mcws, True)) == expected
+    assert list(auto_lowercase_urns(mcws)) == expected
 
 
 @patch("datahub.ingestion.api.source_helpers.telemetry.telemetry_instance.ping")