fix: model package arn unit test

evakravi · evakravi · commit 12678bc20959 · 2023-08-15T15:38:53.000Z
diff --git a/tests/unit/sagemaker/jumpstart/model/test_model.py b/tests/unit/sagemaker/jumpstart/model/test_model.py
@@ -608,7 +608,7 @@ def test_jumpstart_model_package_arn(
             },
         )
 
-        self.assertIn(tag, mock_session.return_value.create_model.call_args[1]["tags"])
+        self.assertIn(tag, mock_session.create_model.call_args[1]["tags"])
 
     @mock.patch("sagemaker.jumpstart.model.is_valid_model_id")
     @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")

Original file line number	Diff line number	Diff line change
`@@ -608,7 +608,7 @@ def test_jumpstart_model_package_arn(`
`608`	`608`	`},`
`609`	`609`	`)`
`610`	`610`
`611`		`- self.assertIn(tag, mock_session.return_value.create_model.call_args[1]["tags"])`
	`611`	`+ self.assertIn(tag, mock_session.create_model.call_args[1]["tags"])`
`612`	`612`
`613`	`613`	`@mock.patch("sagemaker.jumpstart.model.is_valid_model_id")`
`614`	`614`	`@mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")`