fix PR check tests

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."""

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 = {

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')

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

@@ -23,13 +23,43 @@
 )
 
 
+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"]
     return [f"{tag[0]}:{tag[1]}" for tag in tags]
 
 
 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"}

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',