Merge pull request #51 from alphagov/ACW-88/excluding-aliases-with-zero-occurences

Acw 88/excluding aliases with zero occurences

Author: daymohlar

src/models/graph_models.py

@@ -124,6 +124,9 @@ def model_post_init(self, __context: Any) -> None:
         """
         import statistics
 
+        self.alias_imbalance = [stat for stat in self.alias_imbalance if stat.occurrence_count > 0]
+        self.aliases = [alias for alias in self.aliases if alias.occurrence_count > 0]
+
         counts = [stat.occurrence_count for stat in self.alias_imbalance]
         if len(counts) > 1:
             mean = statistics.mean(counts)
@@ -144,3 +147,29 @@ class GraphOutput(BaseModel):
     edges: List[Edge]
     relationships: List[Relationship] = Field(default_factory=list)
     outliers: List[EntityOutlier] = Field(default_factory=list)
+
+    def model_post_init(self, __context: Any) -> None:
+        if not self.outliers:
+            return
+
+        valid_entity_ids = set()
+        valid_node_ids = set()
+
+        for outlier in self.outliers:
+            if len(outlier.alias_imbalance) > 0:
+                valid_entity_ids.add(outlier.entity_id)
+                valid_node_ids.add(outlier.entity_id)
+                for alias in outlier.aliases:
+                    valid_node_ids.add(alias.id)
+
+        self.outliers = [
+            outlier for outlier in self.outliers if outlier.entity_id in valid_entity_ids
+        ]
+
+        self.nodes = [node for node in self.nodes if node.data.id in valid_node_ids]
+
+        self.edges = [
+            edge
+            for edge in self.edges
+            if edge.data.source in valid_node_ids and edge.data.target in valid_node_ids
+        ]

tests/test_graph_loader.py

@@ -55,9 +55,18 @@ def test_local_fixture_loads_independently_of_s3_path():
 def test_graph_output_relationships_contain_required_fields():
     data = load_json_file(FIXTURE_PATH)
     graph = GraphInput.model_validate(data)
-    empty_results: Dict[str, Any] = defaultdict(lambda: defaultdict(list))
 
-    output = build_node_structure(graph.entities, empty_results, graph.relationships)
+    from src.models.graph_models import Occurrence
+
+    mock_occurrence = Occurrence(link="mock", context="mock")
+
+    class MockAliasDict(dict):
+        def get(self, key, default=None):
+            return [mock_occurrence]
+
+    mock_results: Dict[str, Any] = defaultdict(MockAliasDict)
+
+    output = build_node_structure(graph.entities, mock_results, graph.relationships)
     dumped = output.model_dump(exclude_none=True)
 
     assert len(dumped["relationships"]) > 0
@@ -70,9 +79,18 @@ def test_graph_output_relationships_contain_required_fields():
 def test_relationship_edges_reference_existing_nodes():
     data = load_json_file(FIXTURE_PATH)
     graph = GraphInput.model_validate(data)
-    empty_results: Dict[str, Any] = defaultdict(lambda: defaultdict(list))
 
-    output = build_node_structure(graph.entities, empty_results, graph.relationships)
+    from src.models.graph_models import Occurrence
+
+    mock_occurrence = Occurrence(link="mock", context="mock")
+
+    class MockAliasDict(dict):
+        def get(self, key, default=None):
+            return [mock_occurrence]
+
+    mock_results: Dict[str, Any] = defaultdict(MockAliasDict)
+
+    output = build_node_structure(graph.entities, mock_results, graph.relationships)
     dumped = output.model_dump(exclude_none=True)
 
     node_ids = {n["data"]["id"] for n in dumped["nodes"]}

tests/test_graph_validation.py

@@ -44,6 +44,15 @@ def test_graph_output_validation():
     data = {
         "nodes": [{"data": {"id": "e1", "label": "Entity 1", "type": "entity"}}],
         "edges": [{"data": {"source": "e1", "target": "a1", "label": "Alias"}}],
+        "outliers": [
+            {
+                "entity_id": "e1",
+                "entity_label": "Entity 1",
+                "alias_imbalance": [
+                    {"alias_id": "a1", "alias_label": "Alias", "occurrence_count": 1}
+                ],
+            }
+        ],
     }
     validated = GraphOutput.model_validate(data)
     assert len(validated.nodes) == 1