fix test

Author: pintaoz-aws

Diff for: integ/test_codegen.py

@@ -163,7 +163,11 @@ def test_training_and_inference(self):
             content_type="text/csv",
             accept="application/csv",
         )
-        assert invoke_result.body.payload_part
+
+        def deserialise(response):
+            return [res_part for res_part in response.body]
+
+        assert deserialise(invoke_result)
 
     def test_intelligent_defaults(self):
         os.environ["SAGEMAKER_CORE_ADMIN_CONFIG_OVERRIDE"] = (

Diff for: src/sagemaker_core/main/code_injection/codec.py

@@ -206,6 +206,11 @@ def transform(data, shape, object_instance=None) -> dict:
     result = {}
     _shape = SHAPE_DAG[shape]
 
+    if shape == "InvokeEndpointWithResponseStreamOutput":
+        # The ResponseStream shape is different from other shapes
+        # We have to use iterator to iterate it
+        return data
+
     if _shape["type"] in BASIC_TYPES:
         raise ValueError("Unexpected low-level operation model shape")
 

Diff for: src/sagemaker_core/main/utils.py

@@ -500,8 +500,7 @@ def _serialize_dict(value: Dict) -> dict:
     """
     serialized_dict = {}
     for k, v in value.items():
-        serialize_result = serialize(v)
-        if serialize_result is not None:
+        if (serialize_result := serialize(v)) is not None:
             serialized_dict.update({k: serialize_result})
     return serialized_dict
 
@@ -518,8 +517,7 @@ def _serialize_list(value: List) -> list:
     """
     serialized_list = []
     for v in value:
-        serialize_result = serialize(v)
-        if serialize_result is not None:
+        if (serialize_result := serialize(v)) is not None:
             serialized_list.append(serialize_result)
     return serialized_list
 
@@ -536,8 +534,7 @@ def _serialize_shape(value: Any) -> dict:
     """
     serialized_dict = {}
     for k, v in vars(value).items():
-        serialize_result = serialize(v)
-        if serialize_result is not None:
+        if (serialize_result := serialize(v)) is not None:
             key = snake_to_pascal(k) if is_snake_case(k) else k
             serialized_dict.update({key[0].upper() + key[1:]: serialize_result})
     return serialized_dict