1313"""Unit tests for workflow pipeline_experiment_config."""
1414from __future__ import absolute_import
1515
16+ from unittest .mock import Mock
17+
1618from sagemaker .mlops .workflow .pipeline_experiment_config import (
1719 PipelineExperimentConfig , PipelineExperimentConfigProperties
1820)
21+ from sagemaker .mlops .workflow .pipeline import Pipeline , _DEFAULT_EXPERIMENT_CFG
1922from sagemaker .core .workflow .execution_variables import ExecutionVariables
2023
2124
@@ -41,3 +44,46 @@ def test_pipeline_experiment_config_with_execution_variables():
4144def test_pipeline_experiment_config_properties ():
4245 assert PipelineExperimentConfigProperties .EXPERIMENT_NAME .name == "ExperimentName"
4346 assert PipelineExperimentConfigProperties .TRIAL_NAME .name == "TrialName"
47+
48+
49+ # =============================================================================
50+ # Tests for region-based default experiment config handling
51+ # =============================================================================
52+
53+ def _create_mock_session (region : str ) -> Mock :
54+ """Helper to create a mock SageMaker session with specified region."""
55+ mock_session = Mock ()
56+ mock_session .boto_region_name = region
57+ mock_session .boto_session = Mock ()
58+ mock_session .boto_session .client = Mock (return_value = Mock ())
59+ mock_session .local_mode = False
60+ return mock_session
61+
62+
63+ def test_default_config_applied_in_ga_region ():
64+ """Default config applied when nothing provided in GA region."""
65+ mock_session = _create_mock_session ("us-east-1" )
66+ pipeline = Pipeline (name = "test-pipeline" , sagemaker_session = mock_session )
67+ assert pipeline .pipeline_experiment_config == _DEFAULT_EXPERIMENT_CFG
68+
69+
70+ def test_no_default_config_in_non_ga_region ():
71+ """No default config when nothing provided in non-GA region (THE FIX)."""
72+ mock_session = _create_mock_session ("us-gov-west-1" )
73+ pipeline = Pipeline (name = "test-pipeline" , sagemaker_session = mock_session )
74+ assert pipeline .pipeline_experiment_config is None
75+
76+
77+ def test_explicit_none_respected_in_ga_region ():
78+ """Explicit None respected in GA region."""
79+ mock_session = _create_mock_session ("us-east-1" )
80+ pipeline = Pipeline (name = "test-pipeline" , sagemaker_session = mock_session , pipeline_experiment_config = None )
81+ assert pipeline .pipeline_experiment_config is None
82+
83+
84+ def test_custom_config_respected ():
85+ """Custom config respected regardless of region."""
86+ mock_session = _create_mock_session ("us-east-1" )
87+ custom_config = PipelineExperimentConfig ("my-experiment" , "my-trial" )
88+ pipeline = Pipeline (name = "test-pipeline" , sagemaker_session = mock_session , pipeline_experiment_config = custom_config )
89+ assert pipeline .pipeline_experiment_config == custom_config
0 commit comments