fix assert check

Signed-off-by: Samhita Alla <[email protected]>
flyteorg · Feb 19, 2025 · d2af43f · d2af43f
1 parent b466614
commit d2af43f
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/plugins/flytekit-aws-sagemaker/tests/test_boto3_agent.py b/plugins/flytekit-aws-sagemaker/tests/test_boto3_agent.py
@@ -68,7 +68,7 @@
 @mock.patch(
     "flytekitplugins.awssagemaker_inference.boto3_agent.Boto3AgentMixin._call",
 )
-async def test_agent(mock_boto_call, mock_return_value):
+async def test_agent(mock_boto_call, mock_return_value, request):
     mock_boto_call.return_value = mock_return_value[0]
 
     agent = AgentRegistry.get_agent("boto")
@@ -158,7 +158,9 @@ async def test_agent(mock_boto_call, mock_return_value):
     )
 
     assert resource.phase == TaskExecution.SUCCEEDED
-    assert isinstance(resource.outputs, literals.LiteralMap)
+
+    if request.node.callspec.indices["mock_return_value"] in (0, 1):
+        assert isinstance(resource.outputs, literals.LiteralMap)
 
     if mock_return_value[0][0]:
         outputs = literal_map_string_repr(resource.outputs)