Updating unit tests to reflect changes from PR

aviruthen · aviruthen · commit 29498fc898b2 · 2026-01-20T10:19:17.000-08:00
diff --git a/tests/unit/sagemaker/serve/builder/test_model_builder.py b/tests/unit/sagemaker/serve/builder/test_model_builder.py
@@ -316,13 +316,12 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_byoc(
         mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
             mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None
         )
-        mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
+        mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
             (
                 model_data,
                 ENV_VAR_PAIR,
             )
             if model_path == MODEL_PATH
-            and secret_key == mock_secret_key
             and s3_model_data_url == mock_s3_model_data_url
             and sagemaker_session == mock_session
             and image_uri == mock_image_uri
@@ -425,13 +424,12 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_1p_dlc_as_byoc(
         mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
             mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None
         )
-        mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
+        mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
             (
                 model_data,
                 ENV_VAR_PAIR,
             )
             if model_path == MODEL_PATH
-            and secret_key == mock_secret_key
             and s3_model_data_url == mock_s3_model_data_url
             and sagemaker_session == mock_session
             and image_uri == mock_1p_dlc_image_uri
@@ -533,13 +531,12 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_inference_spec(
             if inference_spec == mock_inference_spec and model_server == ModelServer.TORCHSERVE
             else None
         )
-        mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
+        mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
             (
                 model_data,
                 ENV_VAR_PAIR,
             )
             if model_path == MODEL_PATH
-            and secret_key == mock_secret_key
             and s3_model_data_url == mock_s3_model_data_url
             and sagemaker_session == mock_session
             and image_uri == mock_image_uri
@@ -635,13 +632,12 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_model(
         mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
             mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None
         )
-        mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
+        mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
             (
                 model_data,
                 ENV_VAR_PAIR,
             )
             if model_path == MODEL_PATH
-            and secret_key == mock_secret_key
             and s3_model_data_url == mock_s3_model_data_url
             and sagemaker_session == mock_session
             and image_uri == mock_image_uri
@@ -745,13 +741,12 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_xgboost_model(
         mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
             mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None
         )
-        mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
+        mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
             (
                 model_data,
                 ENV_VAR_PAIR,
             )
             if model_path == MODEL_PATH
-            and secret_key == mock_secret_key
             and s3_model_data_url == mock_s3_model_data_url
             and sagemaker_session == mock_session
             and image_uri == mock_image_uri
@@ -973,13 +968,12 @@ def test_build_happy_path_with_localContainer_mode_overwritten_with_sagemaker_mo
             if inference_spec == mock_inference_spec and model_server == ModelServer.TORCHSERVE
             else None
         )
-        mock_sagemaker_endpoint_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
+        mock_sagemaker_endpoint_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
             (
                 model_data,
                 ENV_VAR_PAIR,
             )
             if model_path == MODEL_PATH
-            and secret_key == mock_secret_key
             and s3_model_data_url == mock_s3_model_data_url
             and sagemaker_session == mock_session
             and image_uri == mock_image_uri
@@ -1101,13 +1095,12 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_overwritten_with_local_co
         mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
             mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None
         )
-        mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
+        mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: (  # noqa E501
             (
                 model_data,
                 ENV_VAR_PAIR,
             )
             if model_path == MODEL_PATH
-            and secret_key == mock_secret_key
             and s3_model_data_url == mock_s3_model_data_url
             and sagemaker_session == mock_session
             and image_uri == mock_image_uri
@@ -1129,10 +1122,9 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_overwritten_with_local_co
         )
         mock_lc_mode.prepare.side_effect = lambda: None
         mock_lc_mode.create_server.side_effect = (
-            lambda image_uri, container_timeout_seconds, secret_key, predictor: (
+            lambda image_uri, container_timeout_seconds, predictor, **kwargs: (
                 None
                 if image_uri == mock_image_uri
-                and secret_key == mock_secret_key
                 and container_timeout_seconds == 60
                 else None
             )
diff --git a/tests/unit/sagemaker/serve/detector/test_dependency_manager.py b/tests/unit/sagemaker/serve/detector/test_dependency_manager.py
@@ -99,7 +99,7 @@ def test_capture_dependencies(self, mock_subprocess, mock_file, mock_path):
             call("custom_module==1.2.3\n"),
             call("numpy==4.5\n"),
             call("boto3=1.28.*\n"),
-            call("sagemaker[huggingface]>=2.199\n"),
+            call("sagemaker[huggingface]>=2.257.0,<3.0.0\n"),
             call("other_module@http://some/website.whl\n"),
         ]
         mocked_writes.assert_has_calls(expected_calls)
diff --git a/tests/unit/sagemaker/serve/model_server/tgi/test_tgi_serving.py b/tests/unit/sagemaker/serve/model_server/tgi/test_tgi_serving.py
@@ -54,7 +54,6 @@ def test_tgi_serving_runs_container_non_jumpstart_success(self):
             mock_container_client,
             MOCK_IMAGE,
             MOCK_MODEL_PATH,
-            MOCK_SECRET_KEY,
             MOCK_ENV_VARS,
             False,
         )
@@ -94,7 +93,7 @@ def test_tgi_serving_runs_container_jumpstart_success(self):
 
         # WHEN
         localTgiServing._start_tgi_serving(
-            mock_container_client, MOCK_IMAGE, MOCK_MODEL_PATH, MOCK_SECRET_KEY, MOCK_ENV_VARS, True
+            mock_container_client, MOCK_IMAGE, MOCK_MODEL_PATH, MOCK_ENV_VARS, True
         )
 
         # THEN

Original file line number	Diff line number	Diff line change
`@@ -316,13 +316,12 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_byoc(`
`316`	`316`	`mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (`
`317`	`317`	`mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None`
`318`	`318`	`)`
`319`		`- mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
	`319`	`+ mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
`320`	`320`	`(`
`321`	`321`	`model_data,`
`322`	`322`	`ENV_VAR_PAIR,`
`323`	`323`	`)`
`324`	`324`	`if model_path == MODEL_PATH`
`325`		`- and secret_key == mock_secret_key`
`326`	`325`	`and s3_model_data_url == mock_s3_model_data_url`
`327`	`326`	`and sagemaker_session == mock_session`
`328`	`327`	`and image_uri == mock_image_uri`
`@@ -425,13 +424,12 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_1p_dlc_as_byoc(`
`425`	`424`	`mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (`
`426`	`425`	`mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None`
`427`	`426`	`)`
`428`		`- mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
	`427`	`+ mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
`429`	`428`	`(`
`430`	`429`	`model_data,`
`431`	`430`	`ENV_VAR_PAIR,`
`432`	`431`	`)`
`433`	`432`	`if model_path == MODEL_PATH`
`434`		`- and secret_key == mock_secret_key`
`435`	`433`	`and s3_model_data_url == mock_s3_model_data_url`
`436`	`434`	`and sagemaker_session == mock_session`
`437`	`435`	`and image_uri == mock_1p_dlc_image_uri`
`@@ -533,13 +531,12 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_inference_spec(`
`533`	`531`	`if inference_spec == mock_inference_spec and model_server == ModelServer.TORCHSERVE`
`534`	`532`	`else None`
`535`	`533`	`)`
`536`		`- mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
	`534`	`+ mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
`537`	`535`	`(`
`538`	`536`	`model_data,`
`539`	`537`	`ENV_VAR_PAIR,`
`540`	`538`	`)`
`541`	`539`	`if model_path == MODEL_PATH`
`542`		`- and secret_key == mock_secret_key`
`543`	`540`	`and s3_model_data_url == mock_s3_model_data_url`
`544`	`541`	`and sagemaker_session == mock_session`
`545`	`542`	`and image_uri == mock_image_uri`
`@@ -635,13 +632,12 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_model(`
`635`	`632`	`mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (`
`636`	`633`	`mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None`
`637`	`634`	`)`
`638`		`- mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
	`635`	`+ mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
`639`	`636`	`(`
`640`	`637`	`model_data,`
`641`	`638`	`ENV_VAR_PAIR,`
`642`	`639`	`)`
`643`	`640`	`if model_path == MODEL_PATH`
`644`		`- and secret_key == mock_secret_key`
`645`	`641`	`and s3_model_data_url == mock_s3_model_data_url`
`646`	`642`	`and sagemaker_session == mock_session`
`647`	`643`	`and image_uri == mock_image_uri`
`@@ -745,13 +741,12 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_xgboost_model(`
`745`	`741`	`mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (`
`746`	`742`	`mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None`
`747`	`743`	`)`
`748`		`- mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
	`744`	`+ mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
`749`	`745`	`(`
`750`	`746`	`model_data,`
`751`	`747`	`ENV_VAR_PAIR,`
`752`	`748`	`)`
`753`	`749`	`if model_path == MODEL_PATH`
`754`		`- and secret_key == mock_secret_key`
`755`	`750`	`and s3_model_data_url == mock_s3_model_data_url`
`756`	`751`	`and sagemaker_session == mock_session`
`757`	`752`	`and image_uri == mock_image_uri`
`@@ -973,13 +968,12 @@ def test_build_happy_path_with_localContainer_mode_overwritten_with_sagemaker_mo`
`973`	`968`	`if inference_spec == mock_inference_spec and model_server == ModelServer.TORCHSERVE`
`974`	`969`	`else None`
`975`	`970`	`)`
`976`		`- mock_sagemaker_endpoint_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
	`971`	`+ mock_sagemaker_endpoint_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
`977`	`972`	`(`
`978`	`973`	`model_data,`
`979`	`974`	`ENV_VAR_PAIR,`
`980`	`975`	`)`
`981`	`976`	`if model_path == MODEL_PATH`
`982`		`- and secret_key == mock_secret_key`
`983`	`977`	`and s3_model_data_url == mock_s3_model_data_url`
`984`	`978`	`and sagemaker_session == mock_session`
`985`	`979`	`and image_uri == mock_image_uri`
`@@ -1101,13 +1095,12 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_overwritten_with_local_co`
`1101`	`1095`	`mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (`
`1102`	`1096`	`mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None`
`1103`	`1097`	`)`
`1104`		`- mock_mode.prepare.side_effect = lambda model_path, secret_key, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
	`1098`	`+ mock_mode.prepare.side_effect = lambda model_path, s3_model_data_url, sagemaker_session, image_uri, jumpstart, **kwargs: ( # noqa E501`
`1105`	`1099`	`(`
`1106`	`1100`	`model_data,`
`1107`	`1101`	`ENV_VAR_PAIR,`
`1108`	`1102`	`)`
`1109`	`1103`	`if model_path == MODEL_PATH`
`1110`		`- and secret_key == mock_secret_key`
`1111`	`1104`	`and s3_model_data_url == mock_s3_model_data_url`
`1112`	`1105`	`and sagemaker_session == mock_session`
`1113`	`1106`	`and image_uri == mock_image_uri`
`@@ -1129,10 +1122,9 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_overwritten_with_local_co`
`1129`	`1122`	`)`
`1130`	`1123`	`mock_lc_mode.prepare.side_effect = lambda: None`
`1131`	`1124`	`mock_lc_mode.create_server.side_effect = (`
`1132`		`- lambda image_uri, container_timeout_seconds, secret_key, predictor: (`
	`1125`	`+ lambda image_uri, container_timeout_seconds, predictor, **kwargs: (`
`1133`	`1126`	`None`
`1134`	`1127`	`if image_uri == mock_image_uri`
`1135`		`- and secret_key == mock_secret_key`
`1136`	`1128`	`and container_timeout_seconds == 60`
`1137`	`1129`	`else None`
`1138`	`1130`	`)`
Original file line number	Diff line number	Diff line change
`@@ -99,7 +99,7 @@ def test_capture_dependencies(self, mock_subprocess, mock_file, mock_path):`
`99`	`99`	`call("custom_module==1.2.3\n"),`
`100`	`100`	`call("numpy==4.5\n"),`
`101`	`101`	`call("boto3=1.28.*\n"),`
`102`		`- call("sagemaker[huggingface]>=2.199\n"),`
	`102`	`+ call("sagemaker[huggingface]>=2.257.0,<3.0.0\n"),`
`103`	`103`	`call("other_module@http://some/website.whl\n"),`
`104`	`104`	`]`
`105`	`105`	`mocked_writes.assert_has_calls(expected_calls)`
Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,6 @@ def test_tgi_serving_runs_container_non_jumpstart_success(self):`
`54`	`54`	`mock_container_client,`
`55`	`55`	`MOCK_IMAGE,`
`56`	`56`	`MOCK_MODEL_PATH,`
`57`		`- MOCK_SECRET_KEY,`
`58`	`57`	`MOCK_ENV_VARS,`
`59`	`58`	`False,`
`60`	`59`	`)`
`@@ -94,7 +93,7 @@ def test_tgi_serving_runs_container_jumpstart_success(self):`
`94`	`93`
`95`	`94`	`# WHEN`
`96`	`95`	`localTgiServing._start_tgi_serving(`
`97`		`- mock_container_client, MOCK_IMAGE, MOCK_MODEL_PATH, MOCK_SECRET_KEY, MOCK_ENV_VARS, True`
	`96`	`+ mock_container_client, MOCK_IMAGE, MOCK_MODEL_PATH, MOCK_ENV_VARS, True`
`98`	`97`	`)`
`99`	`98`
`100`	`99`	`# THEN`