Revise agg method and add test

backend/wmg/api/v2.py

@@ -443,11 +443,14 @@ def fill_out_structured_cell_type_agg(cell_type_agg, structured_result, ordering
     n = cell_type_agg["n_cells_cell_type"].values
 
     for i in range(len(tissues)):
+        order = ordering.get((tissues[i], cell_types[i]))
+        if order is None:  # Skip if not found in ordering
+            continue
         structured_result[tissues[i]][cell_types[i]]["aggregated"] = {
             "cell_type_ontology_term_id": cell_types[i],
             "name": ontology_term_label(cell_types[i]),
             "total_count": int(n[i]),
-            "order": ordering.get((tissues[i], cell_types[i]), -1),
+            "order": order,
         }
 
 

tests/unit/backend/wmg/api/test_v2.py

@@ -1,12 +1,14 @@
 import json
 import unittest
+from collections import defaultdict
 from typing import Dict, List
 from unittest.mock import patch
 
+import pandas as pd
 from pytest import approx
 
 from backend.common.census_cube.data.query import MarkerGeneQueryCriteria
-from backend.wmg.api.v2 import find_dimension_id_from_compare
+from backend.wmg.api.v2 import fill_out_structured_cell_type_agg, find_dimension_id_from_compare
 from backend.wmg.server.app import app
 from tests.test_utils import compare_dicts
 from tests.unit.backend.fixtures.environment_setup import EnvironmentSetup
@@ -1425,6 +1427,36 @@ def test__markers_returns_200_and_empty_dictionary_for_bad_celltypes(
             self.assertDictEqual(received, expected)
             self.assertEqual(200, response.status_code)
 
+    @patch("backend.wmg.api.v2.ontology_term_label")
+    def test_fill_out_structured_cell_type_agg(self, ontology_term_label):
+        cell_type_agg = pd.DataFrame(
+            {
+                "tissue_ontology_term_id": ["t1", "t2"],
+                "cell_type_ontology_term_id": ["c1", "c2"],
+                "n_cells_cell_type": [10, 20],
+            }
+        )
+
+        structured_result = defaultdict(lambda: defaultdict(dict))
+
+        ordering = {("t1", "c1"): 1, ("t2", "c2"): 2}
+
+        fill_out_structured_cell_type_agg(cell_type_agg, structured_result, ordering)
+
+        assert structured_result["t1"]["c1"]["aggregated"] == {
+            "cell_type_ontology_term_id": "c1",
+            "name": ontology_term_label("label_c1"),
+            "total_count": 10,
+            "order": 1,
+        }
+
+        assert structured_result["t2"]["c2"]["aggregated"] == {
+            "cell_type_ontology_term_id": "c2",
+            "name": ontology_term_label("label_c2"),
+            "total_count": 20,
+            "order": 2,
+        }
+
 
 # mock the dataset and collection entity data that would otherwise be fetched from the db; in this test
 # we only care that we're building the response correctly from the cube; WMG API integration tests verify