Merge branch 'main' of https://github.com/microsoft/intelligence-toolkit

microsoft · Sep 25, 2024 · 8c8ef64 · 8c8ef64
2 parents 6efa6f6 + ee792c2
commit 8c8ef64
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/toolkit/detect_entity_networks/exposure_report.py b/toolkit/detect_entity_networks/exposure_report.py
@@ -22,6 +22,9 @@ def build_exposure_data(
     if integrated_flags.is_empty():
         return ""
 
+    if isinstance(c_nodes, set):
+        c_nodes = list(c_nodes)
+
     qualified_selected = f"{ENTITY_LABEL}{ATTRIBUTE_VALUE_SEPARATOR}{selected_entity}"
     rdf = integrated_flags
     c_nodes = c_nodes.copy()

diff --git a/toolkit/tests/detect_entity_networks/test_exposure_report.py b/toolkit/tests/detect_entity_networks/test_exposure_report.py
@@ -116,6 +116,22 @@ def test_generation_nodes_inferred_passed(self, graph, flags, c_nodes):
         for ex in expected_nodes:
             assert ex in nodes
 
+    def test_generation_nodes_inferred_c_nodes_set(self, graph, flags, c_nodes) -> None:
+        # c_nodes as set
+        c_nodes = {"ENTITY==A", "ENTITY==C", "ENTITY==D"}
+        graph.add_edge("ENTITY==X", "ENTITY==F")
+        inferred_links = defaultdict(set)
+        inferred_links["ENTITY==X"].add("ENTITY==F")
+        _, _, nodes = build_exposure_data(flags, c_nodes, "X", graph, inferred_links)
+
+        expected_nodes = [
+            {"node": "ENTITY==C", "flags": 2},
+            {"node": "ENTITY==D", "flags": 3},
+            {"node": "ENTITY==X", "flags": 2},
+        ]
+        for ex in expected_nodes:
+            assert ex in nodes
+
     def test_generation_nodes_inferred_path(self, graph, flags, c_nodes) -> None:
         graph.add_edge("ENTITY==X", "ENTITY==F")
         inferred_links = defaultdict(set)