Skip to content

Commit 0bbb83d

Browse files
committed
test: add unit tests for JumpStart network isolation fix
Tests both code paths: 1. _build_for_jumpstart() - verifies enable_network_isolation is applied from init_kwargs and that user-set values are not overridden 2. from_jumpstart_config() - verifies enable_network_isolation is extracted from deploy_kwargs
1 parent 0789d88 commit 0bbb83d

1 file changed

Lines changed: 5 additions & 1 deletion

File tree

sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -381,7 +381,7 @@ def test_from_jumpstart_config_basic(self):
381381
self.assertEqual(mb.model, "test-model")
382382
self.assertEqual(mb.model_version, "1.0.0")
383383

384-
@patch("sagemaker.core.jumpstart.artifacts.kwargs._retrieve_model_deploy_kwargs")
384+
@patch("sagemaker.serve.model_builder._retrieve_model_deploy_kwargs")
385385
def test_from_jumpstart_config_applies_network_isolation(self, mock_deploy_kwargs):
386386
"""Test that enable_network_isolation from deploy kwargs is applied."""
387387
from sagemaker.core.jumpstart.configs import JumpStartConfig
@@ -397,10 +397,14 @@ def test_from_jumpstart_config_applies_network_isolation(self, mock_deploy_kwarg
397397
model_version="1.0.0"
398398
)
399399

400+
mock_session = Mock()
401+
mock_session.boto_region_name = "us-west-2"
402+
400403
mb = ModelBuilder.from_jumpstart_config(
401404
jumpstart_config=js_config,
402405
role_arn="arn:aws:iam::123456789012:role/SageMakerRole",
403406
compute=Compute(instance_type="ml.g5.xlarge"),
407+
sagemaker_session=mock_session,
404408
)
405409

406410
self.assertTrue(mb._enable_network_isolation)

0 commit comments

Comments
 (0)