Fix source retrieval provenance merging

reasoner_pydantic/kgraph.py

@@ -108,7 +108,13 @@ def update(self, other):
                 self.attributes = other.attributes
         if other.sources:
             if self.sources:
+                # We need to do the union of the two sets and then update each
+                # source with the information from the other list
                 self.sources.update(other.sources)
+                for source in self.sources:
+                    for other_source in other.sources:
+                        if source == other_source:
+                            source.update(other_source)
             else:
                 self.sources = other.sources
 

setup.py

@@ -9,7 +9,7 @@
 
 setup(
     name="reasoner-pydantic",
-    version="5.0.3",
+    version="5.0.4",
     author="Abrar Mesbah",
     author_email="amesbah@covar.com",
     url="https://github.com/TranslatorSRI/reasoner-pydantic",

tests/test_update.py

@@ -448,3 +448,86 @@ def test_merge_identical_attributes():
 
     assert ATTRIBUTE_A in node.attributes
     assert len(node.attributes) == 1
+
+
+def test_merge_knowledge_graph_edges():
+    """
+    Test that knowledge graph edges are merged properly
+    """
+
+    message_a = {
+        "knowledge_graph": {
+            "nodes": {},
+            "edges": {
+                "ke0": {
+                    "subject": "kn0",
+                    "object": "kn1",
+                    "predicate": "biolink:ameliorates",
+                    "sources": [
+                        {
+                            "resource_id": "ks0",
+                            "resource_role": "primary_knowledge_source",
+                        },
+                        {
+                            "resource_id": "kp0",
+                            "resource_role": "aggregator_knowledge_source",
+                            "upstream_resource_ids": ["ks0"],
+                        },
+                        {
+                            "resource_id": "ara0",
+                            "resource_role": "aggregator_knowledge_source",
+                            "upstream_resource_ids": ["kp0"],
+                        },
+                    ],
+                    "attributes": [],
+                }
+            },
+        },
+        "results": [],
+    }
+
+    message_b = {
+        "knowledge_graph": {
+            "nodes": {},
+            "edges": {
+                "ke0": {
+                    "subject": "kn0",
+                    "object": "kn1",
+                    "predicate": "biolink:ameliorates",
+                    "sources": [
+                        {
+                            "resource_id": "ks0",
+                            "resource_role": "primary_knowledge_source",
+                        },
+                        {
+                            "resource_id": "kp1",
+                            "resource_role": "aggregator_knowledge_source",
+                            "upstream_resource_ids": ["ks0"],
+                        },
+                        {
+                            "resource_id": "ara0",
+                            "resource_role": "aggregator_knowledge_source",
+                            "upstream_resource_ids": ["kp1"],
+                        },
+                    ],
+                    "attributes": [],
+                }
+            },
+        },
+        "results": [],
+    }
+
+    m = Message()
+
+    m.update(Message.parse_obj(message_a))
+    m.update(Message.parse_obj(message_b))
+
+    edges = m.knowledge_graph.edges
+    assert len(edges) == 1
+    edge = next(iter(edges.values()))
+
+    sources = edge.sources
+    assert len(sources) == 4
+    for source in sources:
+        if source.resource_id == "ara0":
+            assert len(source.upstream_resource_ids) == 2