refactor tests

Signed-off-by: wei-chenglai <[email protected]>

.github/workflows/integration-tests.yaml

@@ -89,7 +89,7 @@ jobs:
       - name: Run initializer_v2 integration tests for Python 3.11+
         if: ${{ matrix.python-version == '3.11' }}
         run: |
-          pip install urllib3 huggingface_hub
+          pip install -r ./cmd/initializer_v2/dataset/requirements.txt
           pip install -U './sdk_v2'
           pytest ./test/integration/initializer_v2
 

cmd/initializer_v2/model/requirements.txt

@@ -1 +1 @@
-huggingface_hub==0.23.4
+huggingface-hub>=0.27.0,<0.28

pkg/initializer_v2/conftest.py

@@ -0,0 +1,23 @@
+import os
+
+import pytest
+
+
+@pytest.fixture
+def mock_env_vars():
+    """Fixture to set and clean up environment variables"""
+    original_env = dict(os.environ)
+
+    def _set_env_vars(**kwargs):
+        for key, value in kwargs.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = str(value)
+        return os.environ
+
+    yield _set_env_vars
+
+    # Cleanup
+    os.environ.clear()
+    os.environ.update(original_env)

pkg/initializer_v2/dataset/main_test.py

@@ -1,31 +1,10 @@
-import os
 from unittest.mock import MagicMock, patch
 
 import pytest
 
 from pkg.initializer_v2.dataset.__main__ import main
 
 
-@pytest.fixture
-def mock_env_vars():
-    """Fixture to set and clean up environment variables"""
-    original_env = dict(os.environ)
-
-    def _set_env_vars(**kwargs):
-        for key, value in kwargs.items():
-            if value is None:
-                os.environ.pop(key, None)
-            else:
-                os.environ[key] = str(value)
-        return os.environ
-
-    yield _set_env_vars
-
-    # Cleanup
-    os.environ.clear()
-    os.environ.update(original_env)
-
-
 @pytest.mark.parametrize(
     "test_name, test_case",
     [
@@ -34,8 +13,6 @@ def _set_env_vars(**kwargs):
             {
                 "storage_uri": "hf://dataset/path",
                 "access_token": "test_token",
-                "mock_config_error": False,
-                "mock_download_error": False,
                 "expected_error": None,
             },
         ),
@@ -44,8 +21,6 @@ def _set_env_vars(**kwargs):
             {
                 "storage_uri": None,
                 "access_token": None,
-                "mock_config_error": False,
-                "mock_download_error": False,
                 "expected_error": Exception,
             },
         ),
@@ -54,28 +29,6 @@ def _set_env_vars(**kwargs):
             {
                 "storage_uri": "invalid://dataset/path",
                 "access_token": None,
-                "mock_config_error": False,
-                "mock_download_error": False,
-                "expected_error": Exception,
-            },
-        ),
-        (
-            "Config loading failure",
-            {
-                "storage_uri": "hf://dataset/path",
-                "access_token": None,
-                "mock_config_error": True,
-                "mock_download_error": False,
-                "expected_error": Exception,
-            },
-        ),
-        (
-            "Dataset download failure",
-            {
-                "storage_uri": "hf://dataset/path/error",
-                "access_token": None,
-                "mock_config_error": False,
-                "mock_download_error": True,
                 "expected_error": Exception,
             },
         ),
@@ -94,10 +47,6 @@ def test_dataset_main(test_name, test_case, mock_env_vars):
 
     # Setup mock HuggingFace instance
     mock_hf_instance = MagicMock()
-    if test_case["mock_config_error"]:
-        mock_hf_instance.load_config.side_effect = Exception
-    if test_case["mock_download_error"]:
-        mock_hf_instance.download_dataset.side_effect = Exception
 
     with patch(
         "pkg.initializer_v2.dataset.__main__.HuggingFace",

pkg/initializer_v2/model/main_test.py

@@ -1,31 +1,10 @@
-import os
 from unittest.mock import MagicMock, patch
 
 import pytest
 
 from pkg.initializer_v2.model.__main__ import main
 
 
-@pytest.fixture
-def mock_env_vars():
-    """Fixture to set and clean up environment variables"""
-    original_env = dict(os.environ)
-
-    def _set_env_vars(**kwargs):
-        for key, value in kwargs.items():
-            if value is None:
-                os.environ.pop(key, None)
-            else:
-                os.environ[key] = str(value)
-        return os.environ
-
-    yield _set_env_vars
-
-    # Cleanup
-    os.environ.clear()
-    os.environ.update(original_env)
-
-
 @pytest.mark.parametrize(
     "test_name, test_case",
     [
@@ -34,8 +13,6 @@ def _set_env_vars(**kwargs):
             {
                 "storage_uri": "hf://model/path",
                 "access_token": "test_token",
-                "mock_config_error": False,
-                "mock_download_error": False,
                 "expected_error": None,
             },
         ),
@@ -44,8 +21,6 @@ def _set_env_vars(**kwargs):
             {
                 "storage_uri": None,
                 "access_token": None,
-                "mock_config_error": False,
-                "mock_download_error": False,
                 "expected_error": Exception,
             },
         ),
@@ -54,28 +29,6 @@ def _set_env_vars(**kwargs):
             {
                 "storage_uri": "invalid://model/path",
                 "access_token": None,
-                "mock_config_error": False,
-                "mock_download_error": False,
-                "expected_error": Exception,
-            },
-        ),
-        (
-            "Config loading failure",
-            {
-                "storage_uri": "hf://model/path",
-                "access_token": None,
-                "mock_config_error": True,
-                "mock_download_error": False,
-                "expected_error": Exception,
-            },
-        ),
-        (
-            "Model download failure",
-            {
-                "storage_uri": "hf://model/path/error",
-                "access_token": None,
-                "mock_config_error": False,
-                "mock_download_error": True,
                 "expected_error": Exception,
             },
         ),
@@ -94,10 +47,6 @@ def test_model_main(test_name, test_case, mock_env_vars):
 
     # Setup mock HuggingFace instance
     mock_hf_instance = MagicMock()
-    if test_case["mock_config_error"]:
-        mock_hf_instance.load_config.side_effect = Exception
-    if test_case["mock_download_error"]:
-        mock_hf_instance.download_model.side_effect = Exception
 
     with patch(
         "pkg.initializer_v2.model.__main__.HuggingFace",

pkg/initializer_v2/utils/utils_test.py

@@ -1,31 +1,9 @@
-import os
-
 import pytest
 from kubeflow.training import HuggingFaceDatasetConfig, HuggingFaceModelInputConfig
 
 import pkg.initializer_v2.utils.utils as utils
 
 
-@pytest.fixture
-def mock_env_vars():
-    """Fixture to set and clean up environment variables"""
-    original_env = dict(os.environ)
-
-    def _set_env_vars(**kwargs):
-        for key, value in kwargs.items():
-            if value is None:
-                os.environ.pop(key, None)
-            else:
-                os.environ[key] = str(value)
-        return os.environ
-
-    yield _set_env_vars
-
-    # Cleanup
-    os.environ.clear()
-    os.environ.update(original_env)
-
-
 @pytest.mark.parametrize(
     "config_class,env_vars,expected",
     [

test/integration/initializer_v2/__init__.py

@@ -1,6 +0,0 @@
-import os
-import sys
-
-sys.path.insert(
-    0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
-)

test/integration/initializer_v2/conftest.py

@@ -0,0 +1,47 @@
+import os
+import shutil
+import tempfile
+
+import pytest
+
+
+@pytest.fixture
+def setup_temp_path(monkeypatch):
+    """Creates temporary directory and patches path constant.
+
+    This fixture:
+    1. Creates a temporary directory
+    2. Allows configuration of path constant
+    3. Handles automatic cleanup after tests
+    4. Restores original environment state
+
+    Args:
+        monkeypatch: pytest fixture for modifying objects
+
+    Returns:
+        function: A configurator that accepts path_var (str) and returns temp_dir path
+
+    Usage:
+        def test_something(setup_temp_path):
+            temp_dir = setup_temp_path("MODEL_PATH")
+            # temp_dir is created and MODEL_PATH is patched
+            # cleanup happens automatically after test
+    """
+    # Setup
+    original_env = dict(os.environ)
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    temp_dir = tempfile.mkdtemp(dir=current_dir)
+
+    def configure_path(path_var: str):
+        """Configure path variable in kubeflow.training"""
+        import kubeflow.training as training
+
+        monkeypatch.setattr(training, path_var, temp_dir)
+        return temp_dir
+
+    yield configure_path
+
+    # Cleanup temp directory after test
+    shutil.rmtree(temp_dir, ignore_errors=True)
+    os.environ.clear()
+    os.environ.update(original_env)

test/integration/initializer_v2/dataset_test.py

@@ -1,10 +1,8 @@
 import os
 import runpy
-import shutil
-import tempfile
+from test.integration.initializer_v2.utils import verify_downloaded_files
 
 import pytest
-from kubeflow.training import DATASET_PATH
 
 import pkg.initializer_v2.utils.utils as utils
 
@@ -13,34 +11,8 @@ class TestDatasetIntegration:
     """Integration tests for dataset initialization"""
 
     @pytest.fixture(autouse=True)
-    def setup_teardown(self, monkeypatch):
-        """Setup and teardown for each test"""
-        # Create temporary directory for dataset downloads
-        current_dir = os.path.dirname(os.path.abspath(__file__))
-        self.temp_dir = tempfile.mkdtemp(dir=current_dir)
-        os.environ[DATASET_PATH] = self.temp_dir
-
-        # Store original environment
-        self.original_env = dict(os.environ)
-
-        # Monkeypatch the constant in the module
-        import kubeflow.training as training
-
-        monkeypatch.setattr(training, "DATASET_PATH", self.temp_dir)
-
-        yield
-
-        # Cleanup
-        shutil.rmtree(self.temp_dir, ignore_errors=True)
-        os.environ.clear()
-        os.environ.update(self.original_env)
-
-    def verify_dataset_files(self, expected_files):
-        """Verify downloaded dataset files"""
-        if expected_files:
-            actual_files = set(os.listdir(self.temp_dir))
-            missing_files = set(expected_files) - actual_files
-            assert not missing_files, f"Missing expected files: {missing_files}"
+    def setup_teardown(self, setup_temp_path):
+        self.temp_dir = setup_temp_path("DATASET_PATH")
 
     @pytest.mark.parametrize(
         "test_name, provider, test_case",
@@ -97,6 +69,6 @@ def test_dataset_download(self, test_name, provider, test_case):
                 )
         else:
             runpy.run_module("pkg.initializer_v2.dataset.__main__", run_name="__main__")
-            self.verify_dataset_files(expected_files)
+            verify_downloaded_files(self.temp_dir, expected_files)
 
         print("Test execution completed")