test_fixes

Author: Gabefire

libs/labelbox/src/labelbox/schema/ontology.py

@@ -569,7 +569,7 @@ class OntologyBuilder:
     def from_dict(cls, dictionary: Dict[str, Any]) -> Dict[str, Any]:
         classifications = []
         for c in dictionary["classifications"]:
-            if "type" in c and c["type"] in PromptResponseClassification.Type:
+            if "type" in c and c["type"] in PromptResponseClassification.Type._value2member_map_.keys():
                 classifications.append(PromptResponseClassification.from_dict(c))
             else:
                 classifications.append(Classification.from_dict(c))

libs/labelbox/src/labelbox/schema/ontology_kind.py

@@ -24,7 +24,7 @@ def get_ontology_kind_validation_error(cls, ontology_kind):
 
     @staticmethod
     def evaluate_ontology_kind_with_media_type(ontology_kind,
-                                               media_type: Union[MediaType, None]) -> MediaType:
+                                               media_type: Union[MediaType, None]) -> Union[MediaType, None]:
 
         if ontology_kind and ontology_kind is OntologyKind.ModelEvaluation:
             if media_type is None: