Update test

Signed-off-by: Kevin Su <[email protected]>

flytekit/core/type_engine.py

@@ -17,7 +17,8 @@
 from google.protobuf.json_format import MessageToDict as _MessageToDict
 from google.protobuf.json_format import ParseDict as _ParseDict
 from google.protobuf.struct_pb2 import Struct
-from marshmallow_enum import LoadDumpOptions
+from marshmallow import fields
+from marshmallow_enum import EnumField, LoadDumpOptions
 from marshmallow_jsonschema import JSONSchema
 
 from flytekit.common.exceptions import user as user_exceptions
@@ -228,7 +229,11 @@ def get_literal_type(self, t: Type[T]) -> LiteralType:
         schema = None
         try:
             s = cast(DataClassJsonMixin, t).schema()
-            s.fields["y"].load_by = LoadDumpOptions.name
+            for _, v in s.fields.items():
+                # marshmallow-jsonschema only supports enums loaded by name.
+                # https://github.com/fuhrysteve/marshmallow-jsonschema/blob/81eada1a0c42ff67de216923968af0a6b54e5dcb/marshmallow_jsonschema/base.py#L228
+                if isinstance(v, EnumField):
+                    v.load_by = LoadDumpOptions.name
             schema = JSONSchema().dump(s)
         except Exception as e:
             logger.warn("failed to extract schema for object %s, (will run schemaless) error: %s", str(t), e)
@@ -780,7 +785,6 @@ def convert_json_schema_to_python_class(schema: dict, schema_name) -> Type[datac
     :param schema: dict representing valid JSON schema
     :param schema_name: dataclass name of return type
     """
-    print(schema)
     attribute_list = []
     for property_key, property_val in schema[schema_name]["properties"].items():
         # Handle list

tests/flytekit/unit/core/test_type_engine.py

@@ -556,9 +556,6 @@ class Datum(object):
         x: int
         y: Color
 
-        def __str__(self):
-            return self.value
-
     lt = TypeEngine.to_literal_type(Datum)
     schema = Datum.schema()
     schema.fields["y"].load_by = LoadDumpOptions.name