Skip to content

Commit f296ac4

Browse files
test: Update coverage improvement test for tests/graphs/abstract_graph_test.py
1 parent 8267e9d commit f296ac4

File tree

1 file changed

+34
-1
lines changed

1 file changed

+34
-1
lines changed

tests/graphs/abstract_graph_test.py

+34-1
Original file line numberDiff line numberDiff line change
@@ -199,4 +199,37 @@ def test_set_common_params(self):
199199
test_params = {"param1": "value1", "param2": "value2"}
200200
graph.set_common_params(test_params)
201201

202-
# Assert that update_config was called on each node with the correct parameters
202+
# Assert that update_config was called on each node with the correct parameters
203+
204+
def test_get_state(self):
205+
"""Test that get_state returns the correct final state with or without a provided key, and raises KeyError for missing keys."""
206+
graph = TestGraph("dummy", {"llm": {"model": "openai/gpt-3.5-turbo", "openai_api_key": "sk-test"}})
207+
# Set a dummy final state
208+
graph.final_state = {"answer": "42", "other": "value"}
209+
# Test without a key returns the entire final_state
210+
state = graph.get_state()
211+
assert state == {"answer": "42", "other": "value"}
212+
# Test with a valid key returns the specific value
213+
answer = graph.get_state("answer")
214+
assert answer == "42"
215+
# Test that a missing key raises a KeyError
216+
with pytest.raises(KeyError):
217+
_ = graph.get_state("nonexistent")
218+
219+
def test_append_node(self):
220+
"""Test that append_node correctly delegates to the graph's append_node method."""
221+
graph = TestGraph("dummy", {"llm": {"model": "openai/gpt-3.5-turbo", "openai_api_key": "sk-test"}})
222+
# Replace the graph object with a mock that has append_node
223+
mock_graph = Mock()
224+
graph.graph = mock_graph
225+
dummy_node = Mock()
226+
graph.append_node(dummy_node)
227+
mock_graph.append_node.assert_called_once_with(dummy_node)
228+
229+
def test_get_execution_info(self):
230+
"""Test that get_execution_info returns the execution info stored in the graph."""
231+
graph = TestGraph("dummy", {"llm": {"model": "openai/gpt-3.5-turbo", "openai_api_key": "sk-test"}})
232+
dummy_info = {"execution": "info", "status": "ok"}
233+
graph.execution_info = dummy_info
234+
info = graph.get_execution_info()
235+
assert info == dummy_info

0 commit comments

Comments
 (0)