File size: 3,084 Bytes
b225a21 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import pytest
from agbenchmark.utils.dependencies.graphs import extract_subgraph_based_on_category
@pytest.fixture
def curriculum_graph():
return {
"edges": [
{"from": "Calculus", "to": "Advanced Calculus"},
{"from": "Algebra", "to": "Calculus"},
{"from": "Biology", "to": "Advanced Biology"},
{"from": "World History", "to": "Modern History"},
],
"nodes": [
{"data": {"category": ["math"]}, "id": "Calculus", "label": "Calculus"},
{
"data": {"category": ["math"]},
"id": "Advanced Calculus",
"label": "Advanced Calculus",
},
{"data": {"category": ["math"]}, "id": "Algebra", "label": "Algebra"},
{"data": {"category": ["science"]}, "id": "Biology", "label": "Biology"},
{
"data": {"category": ["science"]},
"id": "Advanced Biology",
"label": "Advanced Biology",
},
{
"data": {"category": ["history"]},
"id": "World History",
"label": "World History",
},
{
"data": {"category": ["history"]},
"id": "Modern History",
"label": "Modern History",
},
],
}
graph_example = {
"nodes": [
{"id": "A", "data": {"category": []}},
{"id": "B", "data": {"category": []}},
{"id": "C", "data": {"category": ["math"]}},
],
"edges": [{"from": "B", "to": "C"}, {"from": "A", "to": "C"}],
}
def test_dfs_category_math(curriculum_graph):
result_graph = extract_subgraph_based_on_category(curriculum_graph, "math")
# Expected nodes: Algebra, Calculus, Advanced Calculus
# Expected edges: Algebra->Calculus, Calculus->Advanced Calculus
expected_nodes = ["Algebra", "Calculus", "Advanced Calculus"]
expected_edges = [
{"from": "Algebra", "to": "Calculus"},
{"from": "Calculus", "to": "Advanced Calculus"},
]
assert set(node["id"] for node in result_graph["nodes"]) == set(expected_nodes)
assert set((edge["from"], edge["to"]) for edge in result_graph["edges"]) == set(
(edge["from"], edge["to"]) for edge in expected_edges
)
def test_extract_subgraph_math_category():
subgraph = extract_subgraph_based_on_category(graph_example, "math")
assert set(
(node["id"], tuple(node["data"]["category"])) for node in subgraph["nodes"]
) == set(
(node["id"], tuple(node["data"]["category"])) for node in graph_example["nodes"]
)
assert set((edge["from"], edge["to"]) for edge in subgraph["edges"]) == set(
(edge["from"], edge["to"]) for edge in graph_example["edges"]
)
def test_extract_subgraph_non_existent_category():
result_graph = extract_subgraph_based_on_category(graph_example, "toto")
# Asserting that the result graph has no nodes and no edges
assert len(result_graph["nodes"]) == 0
assert len(result_graph["edges"]) == 0
|