diff --git a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py index d6c9f35dbf..ecfa241b73 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py +++ b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py @@ -8,6 +8,8 @@ from dataclasses import dataclass import tempfile +import pytest + from sagemaker.serve.model_builder import ModelBuilder from sagemaker.serve.mode.function_pointers import Mode from sagemaker.serve.utils.types import ModelServer @@ -16,6 +18,26 @@ from sagemaker.core.inference_config import AsyncInferenceConfig from botocore.exceptions import ClientError +TEST_ROLE_ARN = "arn:aws:iam::123456789012:role/SageMakerRole" + + +@pytest.fixture(autouse=True) +def stub_role_resolution(): + """Keep ModelBuilder construction offline. + + ``ModelBuilder.__post_init__`` auto-resolves a serving role when no + ``role_arn`` is given, which calls sts:GetCallerIdentity and the paginated + iam:SimulatePrincipalPolicy. Tests here construct ``ModelBuilder`` without a + role, so unpatched they issue live IAM calls and fail on throttling + (SimulatePrincipalPolicy "Rate exceeded") rather than on the behavior under + test. Tests that patch the resolver themselves still override this. + """ + with patch( + "sagemaker.serve.model_builder.resolve_and_validate_role", + side_effect=lambda provided_role=None, **kwargs: provided_role or TEST_ROLE_ARN, + ): + yield + class TestModelBuilderInit(unittest.TestCase): """Test ModelBuilder initialization.""" diff --git a/sagemaker-train/tests/unit/ai_registry/test_dataset.py b/sagemaker-train/tests/unit/ai_registry/test_dataset.py index 28641567a6..b03e5f31d8 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_dataset.py +++ b/sagemaker-train/tests/unit/ai_registry/test_dataset.py @@ -296,10 +296,27 @@ def mock_exists(path): assert document['DatasetS3Prefix'] == 'path/to/dataset.jsonl' assert document['DatasetS3Bucket'] == 'test-bucket' + @patch('sagemaker.ai_registry.air_hub_entity.AIRHub.get_hub_name', return_value="test-hub") + @patch('sagemaker.ai_registry.dataset._get_default_bucket', return_value="test-bucket") + @patch('sagemaker.train.defaults.TrainDefaults.get_role', return_value="arn:aws:iam::123456789012:role/SageMakerRole") + @patch('sagemaker.train.defaults.TrainDefaults.get_sagemaker_session') + @patch('sagemaker.ai_registry.dataset._get_current_domain_id', return_value=None) + @patch('sagemaker.ai_registry.dataset.Session') @patch('sagemaker.ai_registry.dataset.DataSet._validate_dataset_file') @patch('sagemaker.ai_registry.dataset.DataSet._validate_dataset_format') @patch('sagemaker.ai_registry.dataset.AIRHub') - def test_create_with_local_file(self, mock_air_hub, mock_validate_format, mock_validate_file): + def test_create_with_local_file( + self, + mock_air_hub, + mock_validate_format, + mock_validate_file, + mock_session, + mock_get_domain_id, + mock_get_session, + mock_get_role, + mock_default_bucket, + mock_get_hub_name, + ): mock_air_hub.upload_to_s3.return_value = "s3://bucket/path" mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} mock_air_hub.describe_hub_content.return_value = { diff --git a/sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py b/sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py index da2d8ceca7..c567acf0aa 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py +++ b/sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py @@ -46,7 +46,23 @@ def sample_dataset_file(): class TestDataSetDomainId: """Test domain-id is added to SearchKeywords when available.""" - + + @pytest.fixture(autouse=True) + def stub_account_lookups(self): + """Keep these unit tests offline. + + Uploading a local dataset derives the default bucket, and constructing the entity + derives the hub name — both call STS ``GetCallerIdentity``. ``AIRHub`` is patched on + the base-entity module because ``AIRHubEntity.__init__`` resolves the hub name + through its own import, not the one patched on ``dataset``. + """ + with patch( + 'sagemaker.ai_registry.dataset._get_default_bucket', return_value='test-bucket' + ), patch( + 'sagemaker.ai_registry.air_hub_entity.AIRHub.get_hub_name', return_value='test-hub' + ): + yield + @patch('sagemaker.core.helper.session_helper.Session') @patch('sagemaker.ai_registry.dataset._get_current_domain_id') @patch('sagemaker.ai_registry.dataset.AIRHub') diff --git a/sagemaker-train/tests/unit/ai_registry/test_evaluator.py b/sagemaker-train/tests/unit/ai_registry/test_evaluator.py index eeeedfce58..170528d32c 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_evaluator.py +++ b/sagemaker-train/tests/unit/ai_registry/test_evaluator.py @@ -23,6 +23,9 @@ ) +DUMMY_ROLE = "arn:aws:iam::123456789012:role/SageMakerRole" + + def _keywords_from_import_call(mock_air_hub): """Extract the SearchKeyword strings passed to import_hub_content.""" tags = mock_air_hub.import_hub_content.call_args.kwargs["tags"] @@ -30,6 +33,33 @@ def _keywords_from_import_call(mock_air_hub): class TestEvaluator: + @pytest.fixture(autouse=True) + def stub_aws_resolution(self): + """Keep these unit tests offline. + + ``Evaluator.create`` builds a default ``Session``, auto-detects the Studio domain + ID (STS ``GetCallerIdentity``) and resolves/validates an execution role (IAM + ``SimulatePrincipalPolicy``); constructing the entity also derives the hub name + from the caller's account. Unmocked, those reach real AWS and make the suite slow + and flaky — CI hit ``Throttling: Rate exceeded`` on ``SimulatePrincipalPolicy``. + + ``AIRHub`` is patched on the base-entity module as well as on ``evaluator``, since + ``AIRHubEntity.__init__`` resolves the hub name through its own import. + """ + session = MagicMock() + with patch("sagemaker.ai_registry.evaluator.Session", return_value=session), patch( + "sagemaker.ai_registry.evaluator._get_current_domain_id", return_value=None + ), patch( + "sagemaker.train.defaults.TrainDefaults.get_sagemaker_session", + return_value=session, + ), patch( + "sagemaker.train.defaults.TrainDefaults.get_role", return_value=DUMMY_ROLE + ), patch( + "sagemaker.ai_registry.air_hub_entity.AIRHub.get_hub_name", + return_value="test-hub", + ): + yield + @patch('sagemaker.ai_registry.evaluator.AIRHub') def test_create_with_lambda_arn(self, mock_air_hub): mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} @@ -52,9 +82,10 @@ def test_create_with_lambda_arn(self, mock_air_hub): assert evaluator.method == EvaluatorMethod.LAMBDA mock_air_hub.import_hub_content.assert_called_once() + @patch('sagemaker.ai_registry.evaluator._get_default_bucket', return_value="test-bucket") @patch('sagemaker.ai_registry.evaluator.boto3') @patch('sagemaker.ai_registry.evaluator.AIRHub') - def test_create_with_byoc(self, mock_air_hub, mock_boto3): + def test_create_with_byoc(self, mock_air_hub, mock_boto3, mock_default_bucket): mock_lambda_client = MagicMock() mock_boto3.client.return_value = mock_lambda_client mock_lambda_client.create_function.return_value = {"FunctionArn": "lambda-arn"} diff --git a/sagemaker-train/tests/unit/ai_registry/test_evaluator_domain_id.py b/sagemaker-train/tests/unit/ai_registry/test_evaluator_domain_id.py index f6cdaf44af..6e96f1fc45 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_evaluator_domain_id.py +++ b/sagemaker-train/tests/unit/ai_registry/test_evaluator_domain_id.py @@ -18,18 +18,38 @@ class TestEvaluatorDomainId: """Test domain-id is added to SearchKeywords when available.""" - - @patch('sagemaker.core.helper.session_helper.Session') + + @pytest.fixture(autouse=True) + def stub_aws_resolution(self): + """Keep these unit tests offline. + + ``Evaluator.create`` builds a default ``Session``, resolves/validates an execution + role via IAM ``SimulatePrincipalPolicy``, and derives the hub name from the + caller's account. Unmocked, those reach real AWS and made CI fail with + ``Throttling: Rate exceeded``. + """ + session = Mock() + with patch("sagemaker.ai_registry.evaluator.Session", return_value=session), patch( + "sagemaker.train.defaults.TrainDefaults.get_sagemaker_session", + return_value=session, + ), patch( + "sagemaker.train.defaults.TrainDefaults.get_role", + return_value="arn:aws:iam::123456789012:role/test-role", + ), patch( + "sagemaker.ai_registry.air_hub_entity.AIRHub.get_hub_name", + return_value="test-hub", + ): + yield + @patch('sagemaker.ai_registry.evaluator._get_current_domain_id') @patch('sagemaker.ai_registry.evaluator.AIRHub') def test_domain_id_added_when_available( - self, mock_air_hub, mock_get_domain_id, mock_session + self, mock_air_hub, mock_get_domain_id ): """Test that domain-id is added to tags when available.""" # Setup mocks mock_domain_id = "d-test123456" mock_get_domain_id.return_value = mock_domain_id - mock_session.return_value = Mock() # Mock AIRHub methods mock_air_hub.import_hub_content = Mock() @@ -62,16 +82,14 @@ def test_domain_id_added_when_available( # Verify domain-id is in tags assert any(tag[0] == '@domain' and tag[1] == mock_domain_id for tag in tags) - @patch('sagemaker.core.helper.session_helper.Session') @patch('sagemaker.ai_registry.evaluator._get_current_domain_id') @patch('sagemaker.ai_registry.evaluator.AIRHub') def test_domain_id_not_added_when_unavailable( - self, mock_air_hub, mock_get_domain_id, mock_session + self, mock_air_hub, mock_get_domain_id ): """Test that domain-id is not added when unavailable (non-Studio).""" # Setup mocks - domain_id returns None mock_get_domain_id.return_value = None - mock_session.return_value = Mock() # Mock AIRHub methods mock_air_hub.import_hub_content = Mock() @@ -104,18 +122,16 @@ def test_domain_id_not_added_when_unavailable( # Verify domain-id is NOT in tags assert not any(tag[0] == '@domain' for tag in tags) - @patch('sagemaker.core.helper.session_helper.Session') @patch('sagemaker.ai_registry.evaluator._get_current_domain_id') @patch('sagemaker.ai_registry.evaluator.AIRHub') def test_explicit_domain_id_used_without_auto_detection( - self, mock_air_hub, mock_get_domain_id, mock_session + self, mock_air_hub, mock_get_domain_id ): """An explicit domain_id is tagged and auto-detection is not invoked. Covers the non-Studio case (P467494019) where the domain cannot be inferred and must be supplied by the caller. """ - mock_session.return_value = Mock() mock_air_hub.import_hub_content = Mock() mock_air_hub.describe_hub_content = Mock(return_value={ 'HubContentName': 'test-evaluator',