chore: add tests

Author: evakravi

tests/integ/sagemaker/jumpstart/estimator/test_jumpstart_estimator.py

@@ -19,6 +19,7 @@
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 
 from sagemaker.jumpstart.estimator import JumpStartEstimator
+from sagemaker.jumpstart.types import JumpStartModelSpecs
 import tests
 from tests.integ.sagemaker.jumpstart.constants import (
     ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
@@ -169,6 +170,9 @@ def test_gated_model_training_v2(setup):
         sagemaker_session=get_sm_session(),
     )
 
+    assert isinstance(attached_estimator._internal_config.specs, JumpStartModelSpecs)
+    assert attached_estimator._internal_config.specs.model_id == model_id
+
     # uses ml.g5.2xlarge instance
     predictor = attached_estimator.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

@@ -17,6 +17,7 @@
 
 import pytest
 from sagemaker.enums import EndpointType
+from sagemaker.jumpstart.types import JumpStartModelSpecs
 from sagemaker.predictor import retrieve_default
 
 import tests.integ
@@ -223,6 +224,8 @@ def test_jumpstart_gated_model_inference_component_enabled(setup):
     assert model.model_id == model_id
     assert model.endpoint_name == predictor.endpoint_name
     assert model.inference_component_name == predictor.component_name
+    assert isinstance(model._internal_config.specs, JumpStartModelSpecs)
+    assert model._internal_config.specs.model_id == model_id
 
 
 @mock.patch("sagemaker.jumpstart.cache.JUMPSTART_LOGGER.warning")