FIX UT

Jonathan Makunga · Jonathan Makunga · commit b85e2c32510a · 2024-07-09T15:13:22.000-07:00
diff --git a/tests/unit/sagemaker/serve/utils/test_optimize_utils.py b/tests/unit/sagemaker/serve/utils/test_optimize_utils.py
@@ -338,7 +338,7 @@ def test_with_s3_hf(self, mock_model):
         res_model = _custom_speculative_decoding(mock_model, speculative_decoding_config)
 
         mock_model.add_tags.assert_called_once_with(
-            {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "customer"}
+            {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "custom"}
         )
 
         self.assertEqual(

Original file line number	Diff line number	Diff line change
`@@ -338,7 +338,7 @@ def test_with_s3_hf(self, mock_model):`
`338`	`338`	`res_model = _custom_speculative_decoding(mock_model, speculative_decoding_config)`
`339`	`339`
`340`	`340`	`mock_model.add_tags.assert_called_once_with(`
`341`		`- {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "customer"}`
	`341`	`+ {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "custom"}`
`342`	`342`	`)`
`343`	`343`
`344`	`344`	`self.assertEqual(`