diff --git a/patterns/other/graph_search.py b/patterns/other/graph_search.py index e3c92140..b4d19f00 100644 --- a/patterns/other/graph_search.py +++ b/patterns/other/graph_search.py @@ -78,7 +78,7 @@ def find_shortest_path_bfs(self, start: str, end: str) -> list[str] | None: while len(queue): value = queue.pop(0) - for node in self.graph[value]: + for node in self.graph.get(value, []): if node not in dist_to.keys(): edge_to[node] = value dist_to[node] = dist_to[value] + 1 diff --git a/tests/test_graph_search.py b/tests/test_graph_search.py new file mode 100644 index 00000000..dcbdce1a --- /dev/null +++ b/tests/test_graph_search.py @@ -0,0 +1,16 @@ +from patterns.other.graph_search import GraphSearch + + +def test_find_shortest_path_bfs_handles_leaf_nodes_without_outgoing_edges(): + # Leaf nodes ('C', 'D') have no entry in the graph dict, same as a + # real-world adjacency list that only lists nodes with outgoing edges. + graph = { + "A": ["B"], + "B": ["C", "D"], + } + search = GraphSearch(graph) + + assert search.find_shortest_path_bfs("A", "D") == ["A", "B", "D"] + # Searching for a node that is unreachable must still traverse through + # the leaf nodes without raising KeyError. + assert search.find_shortest_path_bfs("A", "Z") is None