@@ -199,4 +199,37 @@ def test_set_common_params(self):
199
199
test_params = {"param1" : "value1" , "param2" : "value2" }
200
200
graph .set_common_params (test_params )
201
201
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