fix fixture mock

aws · Jun 18, 2024 · addd339 · addd339
1 parent 8bba425
commit addd339
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 10 deletions.
diff --git a/tests/unit/sagemaker/huggingface/test_processing.py b/tests/unit/sagemaker/huggingface/test_processing.py
@@ -19,6 +19,7 @@
 from sagemaker.huggingface.processing import HuggingFaceProcessor
 from sagemaker.fw_utils import UploadedCode
 from sagemaker.session_settings import SessionSettings
+from sagemaker.user_agent import process_studio_metadata_file
 
 from .huggingface_utils import get_full_gpu_image_uri, GPU_INSTANCE_TYPE, REGION
 
@@ -68,11 +69,11 @@ def uploaded_code(
 @pytest.fixture(autouse=True)
 def mock_process_studio_metadata_file(tmp_path):
     studio_file = tmp_path / "resource-metadata.json"
-    studio_file.write_text(json.dumps({"AppType": "studio_type"}))
+    studio_file.write_text(json.dumps({"AppType": "TestAppType"}))
 
     with patch("os.path.exists", return_value=True):
-        with patch("builtins.open", mock_open(read_data=studio_file.read_text())):
-            yield
+        with patch("sagemaker.user_agent.open", mock_open(read_data=studio_file.read_text())):
+            yield process_studio_metadata_file
 
 
 @patch("sagemaker.utils._botocore_resolver")

diff --git a/tests/unit/sagemaker/workflow/test_processing_step.py b/tests/unit/sagemaker/workflow/test_processing_step.py
@@ -56,6 +56,7 @@
 from sagemaker.network import NetworkConfig
 from sagemaker.pytorch.estimator import PyTorch
 from sagemaker import utils, Model
+from sagemaker.user_agent import process_studio_metadata_file
 
 from sagemaker.clarify import (
     SageMakerClarifyProcessor,
@@ -247,14 +248,13 @@ def network_config():
     )
 
 
-@pytest.fixture(autouse=True)
 def mock_process_studio_metadata_file(tmp_path):
     studio_file = tmp_path / "resource-metadata.json"
-    studio_file.write_text(json.dumps({"AppType": "studio_type"}))
+    studio_file.write_text(json.dumps({"AppType": "TestAppType"}))
 
     with patch("os.path.exists", return_value=True):
-        with patch("builtins.open", mock_open(read_data=studio_file.read_text())):
-            yield
+        with patch("sagemaker.user_agent.open", mock_open(read_data=studio_file.read_text())):
+            yield process_studio_metadata_file
 
 
 @pytest.mark.parametrize(

diff --git a/tests/unit/test_processing.py b/tests/unit/test_processing.py
@@ -47,6 +47,7 @@
 from sagemaker.workflow.pipeline_context import PipelineSession, _PipelineConfig
 from sagemaker.workflow.functions import Join
 from sagemaker.workflow.execution_variables import ExecutionVariables
+from sagemaker.user_agent import process_studio_metadata_file
 from tests.unit import SAGEMAKER_CONFIG_PROCESSING_JOB
 
 BUCKET_NAME = "mybucket"
@@ -137,11 +138,11 @@ def uploaded_code(
 @pytest.fixture(autouse=True)
 def mock_process_studio_metadata_file(tmp_path):
     studio_file = tmp_path / "resource-metadata.json"
-    studio_file.write_text(json.dumps({"AppType": "studio_type"}))
+    studio_file.write_text(json.dumps({"AppType": "TestAppType"}))
 
     with patch("os.path.exists", return_value=True):
-        with patch("builtins.open", mock_open(read_data=studio_file.read_text())):
-            yield
+        with patch("sagemaker.user_agent.open", mock_open(read_data=studio_file.read_text())):
+            yield process_studio_metadata_file
 
 
 @patch("sagemaker.utils._botocore_resolver")