test get headers from object

monarch-initiative · Jan 8, 2024 · 4f94bce · 4f94bce
1 parent b4d8a05
commit 4f94bce
Show file tree

Hide file tree

Showing 2 changed files with 160 additions and 7 deletions.
diff --git a/backend/src/monarch_py/utils/utils.py b/backend/src/monarch_py/utils/utils.py
@@ -76,6 +76,7 @@ def set_log_level(log_level: str):
 
 ### URL fetching methods ###
 
+
 def get_links_for_field(field: List[str]) -> List[ExpandedCurie]:
     # TODO should be able to remove curie.replace("PMID", "PUBMED")) since the converter should handle prefix synonyms
     return [ExpandedCurie(id=curie, url=converter.expand(curie.replace("PMID", "PUBMED"))) for curie in field]
@@ -100,10 +101,16 @@ def get_provided_by_link(provided_by: str) -> ExpandedCurie:
 
 def get_headers_from_obj(obj: ConfiguredBaseModel) -> list:
     """Return a list of headers from a pydantic model."""
-    schema = type(obj).schema()
-    definitions = schema["definitions"]
-    this_ref = schema["properties"]["items"]["items"]["$ref"].split("/")[-1]
-    headers = definitions[this_ref]["properties"].keys()
+    if isinstance(obj, Entity):
+        headers = obj.dict().keys()
+    elif isinstance(obj, (AssociationCountList, HistoPheno, Results)):
+        if obj.items:
+            headers = obj.items[0].dict().keys()
+        else:
+            schema = type(obj).schema()
+            definitions = schema["definitions"]
+            this_ref = schema["properties"]["items"]["items"]["$ref"].split("/")[-1]
+            headers = definitions[this_ref]["properties"].keys()
     return list(headers)
 
 

diff --git a/backend/tests/unit/test_utils.py b/backend/tests/unit/test_utils.py
@@ -3,10 +3,11 @@
 
 from monarch_py.datamodels.model import (
     AssociationCountList,
-    Entity,
-    HistoPheno,
     AssociationResults,
-    ExpandedCurie
+    HistoPheno,
+    Node,
+    SearchResults,
+    ExpandedCurie,
 )
 from monarch_py.utils.utils import (
     compare_dicts,
@@ -22,6 +23,7 @@
 
 ### Test basic utility methods ###
 
+
 def test_strip_json():
     doc = {"a": 1, "b": 2, "c": 3}
     assert strip_json(doc, "a", "b") == {"c": 3}
@@ -69,8 +71,10 @@ def test_compare_dicts_false():
 def test_dict_diff(d1, d2, expected):
     assert dict_diff(d1, d2) == expected
 
+
 ### Test URL fetching methods ###
 
+
 def test_get_links_for_field():
     field = ["PMID:123", "PMID:456"]
     assert get_links_for_field(field) == [
@@ -87,3 +91,145 @@ def test_get_provided_by_link():
 
 ### Test output conversion methods ###
 # TODO
+
+
+@pytest.mark.parametrize(
+    "obj, expected",
+    [
+        (
+            "node",
+            [
+                "id",
+                "category",
+                "name",
+                "full_name",
+                "deprecated",
+                "description",
+                "xref",
+                "provided_by",
+                "in_taxon",
+                "in_taxon_label",
+                "symbol",
+                "synonym",
+                "uri",
+                "inheritance",
+                "causal_gene",
+                "causes_disease",
+                "mappings",
+                "external_links",
+                "provided_by_link",
+                "association_counts",
+                "node_hierarchy",
+            ],
+        ),
+        (
+            "search",
+            [
+                "id",
+                "category",
+                "name",
+                "full_name",
+                "deprecated",
+                "description",
+                "xref",
+                "provided_by",
+                "in_taxon",
+                "in_taxon_label",
+                "symbol",
+                "synonym",
+                "uri",
+                "highlight",
+                "score",
+            ],
+        ),
+        (
+            "histopheno",
+            ["label", "count", "id"],
+        ),
+        (
+            "associations",
+            [
+                "id",
+                "category",
+                "subject",
+                "original_subject",
+                "subject_namespace",
+                "subject_category",
+                "subject_closure",
+                "subject_label",
+                "subject_closure_label",
+                "subject_taxon",
+                "subject_taxon_label",
+                "predicate",
+                "object",
+                "original_object",
+                "object_namespace",
+                "object_category",
+                "object_closure",
+                "object_label",
+                "object_closure_label",
+                "object_taxon",
+                "object_taxon_label",
+                "primary_knowledge_source",
+                "aggregator_knowledge_source",
+                "negated",
+                "pathway",
+                "evidence_count",
+                "has_evidence",
+                "has_evidence_links",
+                "grouping_key",
+                "provided_by",
+                "provided_by_link",
+                "publications",
+                "publications_links",
+                "qualifiers",
+                "frequency_qualifier",
+                "onset_qualifier",
+                "sex_qualifier",
+                "stage_qualifier",
+                "qualifiers_label",
+                "qualifiers_namespace",
+                "qualifiers_category",
+                "qualifiers_closure",
+                "qualifiers_closure_label",
+                "frequency_qualifier_label",
+                "frequency_qualifier_namespace",
+                "frequency_qualifier_category",
+                "frequency_qualifier_closure",
+                "frequency_qualifier_closure_label",
+                "onset_qualifier_label",
+                "onset_qualifier_namespace",
+                "onset_qualifier_category",
+                "onset_qualifier_closure",
+                "onset_qualifier_closure_label",
+                "sex_qualifier_label",
+                "sex_qualifier_namespace",
+                "sex_qualifier_category",
+                "sex_qualifier_closure",
+                "sex_qualifier_closure_label",
+                "stage_qualifier_label",
+                "stage_qualifier_namespace",
+                "stage_qualifier_category",
+                "stage_qualifier_closure",
+                "stage_qualifier_closure_label",
+            ],
+        ),
+    ],
+)
+def test_get_headers_from_obj(obj, expected, request):
+    if obj == "node":
+        node = request.getfixturevalue(obj)
+        obj = Node(**node)
+    elif obj == "search":
+        search = request.getfixturevalue(obj)
+        obj = SearchResults(**search)
+    elif obj == "histopheno":
+        histopheno = request.getfixturevalue(obj)
+        obj = HistoPheno(**histopheno)
+    elif obj == "associations":
+        associations = request.getfixturevalue(obj)
+        obj = AssociationResults(**associations)
+
+    headers = get_headers_from_obj(obj)
+    print(headers)
+    assert headers == expected