fix(tests): correct mock paths for lazy imports in hub factory and transformation tests

Fixed 20 failing tests across test_dataset_hub_factory.py and
test_dataset_transformation.py. Root cause: lazy imports inside methods
meant mock patches on the hub_factory module couldn't find the target
attributes. Patched at source modules instead (dataset.DataSet,
dataset_transformation.DatasetTransformation, air_hub.AIRHub). Also
fixed test_unsupported_format_raises to use DPO (not in inbound
converters) instead of converse (which is), and fixed detect_format
tests to mock DatasetFormatDetector at its source module.

---
X-AI-Prompt: fix failing unit tests for dataset_hub_factory and dataset_transformation
X-AI-Tool: Kiro

Author: gwang111

sagemaker-train/src/sagemaker/ai_registry/dataset_hub_factory.py

@@ -17,14 +17,15 @@
 import logging
 import os
 import tempfile
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, Union
 
-from sagemaker.ai_registry.dataset import DataSet
 from sagemaker.ai_registry.dataset_transformation import DatasetFormat
 from sagemaker.ai_registry.dataset_utils import CustomizationTechnique
-from sagemaker.ai_registry.air_hub import AIRHub
 from sagemaker.core.helper.session_helper import Session
 
+if TYPE_CHECKING:
+    from sagemaker.ai_registry.dataset import DataSet
+
 logger = logging.getLogger(__name__)
 
 
@@ -52,9 +53,10 @@ def _resolve_dataset(
         Returns:
             A hydrated DataSet instance.
         """
-        if isinstance(dataset, DataSet):
-            return dataset
-        return DataSet.get(name=dataset, sagemaker_session=sagemaker_session)
+        if isinstance(dataset, str):
+            from sagemaker.ai_registry.dataset import DataSet
+            return DataSet.get(name=dataset, sagemaker_session=sagemaker_session)
+        return dataset
 
     @classmethod
     def _download_to_local(cls, s3_uri: str) -> str:
@@ -69,6 +71,7 @@ def _download_to_local(cls, s3_uri: str) -> str:
         suffix = os.path.splitext(s3_uri)[-1] or ".jsonl"
         tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
         tmp.close()
+        from sagemaker.ai_registry.air_hub import AIRHub
         AIRHub.download_from_s3(s3_uri, tmp.name)
         return tmp.name
 
@@ -106,6 +109,7 @@ def transform_dataset(
             ValueError: If neither or both of source and dataset are provided.
         """
         from sagemaker.ai_registry.dataset_transformation import DatasetTransformation
+        from sagemaker.ai_registry.dataset import DataSet
 
         if (source is None) == (dataset is None):
             raise ValueError(

sagemaker-train/tests/unit/ai_registry/test_dataset_hub_factory.py

@@ -11,31 +11,27 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-import json
 import os
 import pytest
-from unittest.mock import Mock, patch, MagicMock
+from unittest.mock import Mock, patch
 
 from sagemaker.ai_registry.dataset_hub_factory import DataSetHubFactory
-from sagemaker.ai_registry.dataset import DataSet
 from sagemaker.ai_registry.dataset_transformation import DatasetFormat
 from sagemaker.ai_registry.dataset_utils import CustomizationTechnique
 from sagemaker.ai_registry.air_constants import HubContentStatus
 
 
 def _make_dataset(**overrides):
-    """Helper to create a DataSet instance with sensible defaults."""
-    defaults = dict(
-        name="test-ds",
-        arn="arn:aws:sagemaker:us-east-1:123456789012:hub-content/test-ds",
-        version="1.0.0",
-        status=HubContentStatus.AVAILABLE,
-        source="s3://bucket/datasets/test-ds/data.jsonl",
-        description="test",
-        customization_technique=CustomizationTechnique.SFT,
-    )
-    defaults.update(overrides)
-    return DataSet(**defaults)
+    """Helper to create a mock DataSet with sensible defaults."""
+    ds = Mock()
+    ds.name = overrides.get("name", "test-ds")
+    ds.arn = overrides.get("arn", "arn:aws:sagemaker:us-east-1:123456789012:hub-content/test-ds")
+    ds.version = overrides.get("version", "1.0.0")
+    ds.status = overrides.get("status", HubContentStatus.AVAILABLE)
+    ds.source = overrides.get("source", "s3://bucket/datasets/test-ds/data.jsonl")
+    ds.description = overrides.get("description", "test")
+    ds.customization_technique = overrides.get("customization_technique", CustomizationTechnique.SFT)
+    return ds
 
 
 class TestResolveDataset:
@@ -44,14 +40,14 @@ def test_resolve_with_dataset_instance(self):
         result = DataSetHubFactory._resolve_dataset(ds)
         assert result is ds
 
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSet.get")
+    @patch("sagemaker.ai_registry.dataset.DataSet.get")
     def test_resolve_with_name_string(self, mock_get):
         mock_get.return_value = _make_dataset()
         result = DataSetHubFactory._resolve_dataset("my-dataset")
         mock_get.assert_called_once_with(name="my-dataset", sagemaker_session=None)
         assert result.name == "test-ds"
 
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSet.get")
+    @patch("sagemaker.ai_registry.dataset.DataSet.get")
     def test_resolve_with_session(self, mock_get):
         mock_session = Mock()
         mock_get.return_value = _make_dataset()
@@ -60,7 +56,7 @@ def test_resolve_with_session(self, mock_get):
 
 
 class TestDownloadToLocal:
-    @patch("sagemaker.ai_registry.dataset_hub_factory.AIRHub.download_from_s3")
+    @patch("sagemaker.ai_registry.air_hub.AIRHub.download_from_s3")
     def test_download_creates_temp_file(self, mock_download):
         path = DataSetHubFactory._download_to_local("s3://bucket/data.jsonl")
         try:
@@ -70,7 +66,7 @@ def test_download_creates_temp_file(self, mock_download):
             if os.path.exists(path):
                 os.unlink(path)
 
-    @patch("sagemaker.ai_registry.dataset_hub_factory.AIRHub.download_from_s3")
+    @patch("sagemaker.ai_registry.air_hub.AIRHub.download_from_s3")
     def test_download_default_suffix(self, mock_download):
         path = DataSetHubFactory._download_to_local("s3://bucket/data")
         try:
@@ -81,11 +77,15 @@ def test_download_default_suffix(self, mock_download):
 
 
 class TestTransformDatasetValidation:
-    def test_neither_source_nor_dataset_raises(self):
+    @patch("sagemaker.ai_registry.dataset_transformation.DatasetTransformation")
+    @patch("sagemaker.ai_registry.dataset.DataSet")
+    def test_neither_source_nor_dataset_raises(self, mock_ds, mock_transform):
         with pytest.raises(ValueError, match="Exactly one of"):
             DataSetHubFactory.transform_dataset(name="test", target_format=DatasetFormat.GENQA)
 
-    def test_both_source_and_dataset_raises(self):
+    @patch("sagemaker.ai_registry.dataset_transformation.DatasetTransformation")
+    @patch("sagemaker.ai_registry.dataset.DataSet")
+    def test_both_source_and_dataset_raises(self, mock_ds, mock_transform):
         with pytest.raises(ValueError, match="Exactly one of"):
             DataSetHubFactory.transform_dataset(
                 name="test",
@@ -96,12 +96,12 @@ def test_both_source_and_dataset_raises(self):
 
 
 class TestTransformDatasetFromSource:
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSet.create")
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DatasetTransformation")
-    def test_local_file_source(self, mock_transformation_cls, mock_create):
+    @patch("sagemaker.ai_registry.dataset.DataSet")
+    @patch("sagemaker.ai_registry.dataset_transformation.DatasetTransformation")
+    def test_local_file_source(self, mock_transformation_cls, mock_dataset_cls):
         mock_transformation_cls.detect_format.return_value = DatasetFormat.OPENAI_CHAT
         mock_transformation_cls.transform_file.return_value = "/tmp/transformed.jsonl"
-        mock_create.return_value = _make_dataset()
+        mock_dataset_cls.create.return_value = _make_dataset()
 
         result = DataSetHubFactory.transform_dataset(
             name="new-ds",
@@ -116,26 +116,26 @@ def test_local_file_source(self, mock_transformation_cls, mock_create):
             source_format=DatasetFormat.OPENAI_CHAT,
             target_format=DatasetFormat.GENQA,
         )
-        mock_create.assert_called_once_with(
+        mock_dataset_cls.create.assert_called_once_with(
             name="new-ds",
             source="/tmp/transformed.jsonl",
             customization_technique=CustomizationTechnique.SFT,
             sagemaker_session=None,
         )
         assert result.name == "test-ds"
 
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSet.create")
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DatasetTransformation")
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSetHubFactory._download_to_local")
+    @patch("sagemaker.ai_registry.dataset.DataSet")
+    @patch("sagemaker.ai_registry.dataset_transformation.DatasetTransformation")
+    @patch.object(DataSetHubFactory, "_download_to_local")
     @patch("os.path.exists", return_value=True)
     @patch("os.remove")
     def test_s3_source_downloads_and_cleans_up(
-        self, mock_remove, mock_exists, mock_download, mock_transformation_cls, mock_create
+        self, mock_remove, mock_exists, mock_download, mock_transformation_cls, mock_dataset_cls
     ):
         mock_download.return_value = "/tmp/downloaded.jsonl"
         mock_transformation_cls.detect_format.return_value = DatasetFormat.HF_PROMPT_COMPLETION
         mock_transformation_cls.transform_file.return_value = "/tmp/transformed.jsonl"
-        mock_create.return_value = _make_dataset()
+        mock_dataset_cls.create.return_value = _make_dataset()
 
         result = DataSetHubFactory.transform_dataset(
             name="new-ds",
@@ -147,12 +147,13 @@ def test_s3_source_downloads_and_cleans_up(
         mock_transformation_cls.detect_format.assert_called_once_with("/tmp/downloaded.jsonl")
         mock_remove.assert_called_once_with("/tmp/downloaded.jsonl")
 
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DatasetTransformation")
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSetHubFactory._download_to_local")
+    @patch("sagemaker.ai_registry.dataset_transformation.DatasetTransformation")
+    @patch("sagemaker.ai_registry.dataset.DataSet")
+    @patch.object(DataSetHubFactory, "_download_to_local")
     @patch("os.path.exists", return_value=True)
     @patch("os.remove")
     def test_s3_source_cleans_up_on_error(
-        self, mock_remove, mock_exists, mock_download, mock_transformation_cls
+        self, mock_remove, mock_exists, mock_download, mock_ds, mock_transformation_cls
     ):
         mock_download.return_value = "/tmp/downloaded.jsonl"
         mock_transformation_cls.detect_format.side_effect = ValueError("bad format")
@@ -168,19 +169,19 @@ def test_s3_source_cleans_up_on_error(
 
 
 class TestTransformDatasetFromExisting:
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSet.get")
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DatasetTransformation")
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSetHubFactory._download_to_local")
+    @patch("sagemaker.ai_registry.dataset.DataSet")
+    @patch("sagemaker.ai_registry.dataset_transformation.DatasetTransformation")
+    @patch.object(DataSetHubFactory, "_download_to_local")
     @patch("os.path.exists", return_value=True)
     @patch("os.remove")
     def test_existing_dataset_instance(
-        self, mock_remove, mock_exists, mock_download, mock_transformation_cls, mock_get
+        self, mock_remove, mock_exists, mock_download, mock_transformation_cls, mock_dataset_cls
     ):
         ds = _make_dataset()
         mock_download.return_value = "/tmp/downloaded.jsonl"
         mock_transformation_cls.detect_format.return_value = DatasetFormat.VERL
         mock_transformation_cls.transform_file.return_value = "/tmp/transformed.jsonl"
-        mock_get.return_value = _make_dataset(version="2.0.0")
+        mock_dataset_cls.get.return_value = _make_dataset(version="2.0.0")
 
         result = DataSetHubFactory.transform_dataset(
             name="test-ds",
@@ -197,10 +198,10 @@ def test_existing_dataset_instance(
         )
         assert result.version == "2.0.0"
 
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSet.get")
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DatasetTransformation")
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSetHubFactory._download_to_local")
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSetHubFactory._resolve_dataset")
+    @patch("sagemaker.ai_registry.dataset.DataSet")
+    @patch("sagemaker.ai_registry.dataset_transformation.DatasetTransformation")
+    @patch.object(DataSetHubFactory, "_download_to_local")
+    @patch.object(DataSetHubFactory, "_resolve_dataset")
     @patch("os.path.exists", return_value=True)
     @patch("os.remove")
     def test_existing_dataset_by_name(
@@ -210,14 +211,14 @@ def test_existing_dataset_by_name(
         mock_resolve,
         mock_download,
         mock_transformation_cls,
-        mock_get,
+        mock_dataset_cls,
     ):
         resolved_ds = _make_dataset()
         mock_resolve.return_value = resolved_ds
         mock_download.return_value = "/tmp/downloaded.jsonl"
         mock_transformation_cls.detect_format.return_value = DatasetFormat.OPENAI_CHAT
         mock_transformation_cls.transform_file.return_value = "/tmp/transformed.jsonl"
-        mock_get.return_value = _make_dataset(version="2.0.0")
+        mock_dataset_cls.get.return_value = _make_dataset(version="2.0.0")
 
         result = DataSetHubFactory.transform_dataset(
             name="test-ds",
@@ -228,12 +229,13 @@ def test_existing_dataset_by_name(
         mock_resolve.assert_called_once_with("test-ds", None)
         mock_download.assert_called_once_with(resolved_ds.source)
 
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DatasetTransformation")
-    @patch("sagemaker.ai_registry.dataset_hub_factory.DataSetHubFactory._download_to_local")
+    @patch("sagemaker.ai_registry.dataset_transformation.DatasetTransformation")
+    @patch("sagemaker.ai_registry.dataset.DataSet")
+    @patch.object(DataSetHubFactory, "_download_to_local")
     @patch("os.path.exists", return_value=True)
     @patch("os.remove")
     def test_existing_dataset_cleans_up_on_error(
-        self, mock_remove, mock_exists, mock_download, mock_transformation_cls
+        self, mock_remove, mock_exists, mock_download, mock_ds, mock_transformation_cls
     ):
         ds = _make_dataset()
         mock_download.return_value = "/tmp/downloaded.jsonl"

sagemaker-train/tests/unit/ai_registry/test_dataset_transformation.py

@@ -332,7 +332,7 @@ def test_openai_dispatch(self):
 
     def test_unsupported_format_raises(self):
         with pytest.raises(ValueError, match="Unsupported dataset format"):
-            convert_to_genqa({}, DatasetFormat("converse"))
+            convert_to_genqa({}, DatasetFormat.DPO)
 
 
 class TestConvertFromGenqa:
@@ -347,21 +347,21 @@ def test_unimplemented_raises(self):
 
 
 class TestDatasetTransformationDetectFormat:
-    @patch("sagemaker.ai_registry.dataset_transformation.DatasetFormatDetector")
-    def test_detect_known_format(self, mock_detector_cls):
-        mock_detector_cls.detect_format.return_value = "openai_chat"
+    @patch("sagemaker.ai_registry.dataset_format_detector.DatasetFormatDetector.detect_format")
+    def test_detect_known_format(self, mock_detect):
+        mock_detect.return_value = "openai_chat"
         result = DatasetTransformation.detect_format("/path/to/file.jsonl")
         assert result == DatasetFormat.OPENAI_CHAT
 
-    @patch("sagemaker.ai_registry.dataset_transformation.DatasetFormatDetector")
-    def test_detect_none_raises(self, mock_detector_cls):
-        mock_detector_cls.detect_format.return_value = None
+    @patch("sagemaker.ai_registry.dataset_format_detector.DatasetFormatDetector.detect_format")
+    def test_detect_none_raises(self, mock_detect):
+        mock_detect.return_value = None
         with pytest.raises(ValueError, match="Unable to detect dataset format"):
             DatasetTransformation.detect_format("/path/to/file.jsonl")
 
-    @patch("sagemaker.ai_registry.dataset_transformation.DatasetFormatDetector")
-    def test_detect_unsupported_format_raises(self, mock_detector_cls):
-        mock_detector_cls.detect_format.return_value = "rft"
+    @patch("sagemaker.ai_registry.dataset_format_detector.DatasetFormatDetector.detect_format")
+    def test_detect_unsupported_format_raises(self, mock_detect):
+        mock_detect.return_value = "rft"
         with pytest.raises(ValueError, match="not a supported transformation format"):
             DatasetTransformation.detect_format("/path/to/file.jsonl")