From d2af43fe45e53300d50731c8d9dddab825e56dca Mon Sep 17 00:00:00 2001 From: Samhita Alla Date: Wed, 19 Feb 2025 23:21:17 +0530 Subject: [PATCH] fix assert check Signed-off-by: Samhita Alla --- plugins/flytekit-aws-sagemaker/tests/test_boto3_agent.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/plugins/flytekit-aws-sagemaker/tests/test_boto3_agent.py b/plugins/flytekit-aws-sagemaker/tests/test_boto3_agent.py index 98e21a362c..f57194dc7c 100644 --- 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)