forked from HAL9000/cleveragents-core
783 lines
28 KiB
Python
783 lines
28 KiB
Python
"""
|
|
BDD step definitions for comprehensive state management coverage testing.
|
|
"""
|
|
|
|
import json
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
from behave import given, then, when
|
|
from behave.runner import Context
|
|
|
|
from cleveragents.langgraph.state import (
|
|
GraphState,
|
|
StateManager,
|
|
StateUpdateMode,
|
|
)
|
|
|
|
|
|
@given("the state management system is available")
|
|
def step_state_management_available(context: Context):
|
|
"""Initialize state management system for testing."""
|
|
context.state_instances = {}
|
|
context.state_managers = {}
|
|
context.test_data = {}
|
|
context.results = {}
|
|
|
|
|
|
@given("I have a GraphState instance")
|
|
def step_create_graph_state(context: Context):
|
|
"""Create a GraphState instance."""
|
|
context.state_instances["main"] = GraphState(
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
metadata={"version": "1.0"},
|
|
current_node="start",
|
|
execution_count=5,
|
|
error=None,
|
|
)
|
|
|
|
|
|
@given("I have updates with replace mode")
|
|
def step_create_replace_mode_updates(context: Context):
|
|
"""Create updates for replace mode testing."""
|
|
context.test_data["replace_updates"] = {
|
|
"current_node": "new_node",
|
|
"execution_count": 10,
|
|
"invalid_attribute": "should_not_be_set",
|
|
}
|
|
|
|
|
|
@when("I update the state with REPLACE mode")
|
|
def step_update_state_replace_mode(context: Context):
|
|
"""Update state using REPLACE mode."""
|
|
state = context.state_instances["main"]
|
|
updates = context.test_data["replace_updates"]
|
|
state.update(updates, StateUpdateMode.REPLACE)
|
|
|
|
|
|
@then("the state should be replaced completely")
|
|
def step_verify_state_replaced(context: Context):
|
|
"""Verify state was replaced completely."""
|
|
state = context.state_instances["main"]
|
|
assert state.current_node == "new_node"
|
|
assert state.execution_count == 10
|
|
|
|
|
|
@then("only valid attributes should be updated")
|
|
def step_verify_only_valid_attributes(context: Context):
|
|
"""Verify only valid attributes are updated."""
|
|
state = context.state_instances["main"]
|
|
# Invalid attributes should not be set
|
|
assert not hasattr(state, "invalid_attribute")
|
|
|
|
|
|
@given("I have a GraphState with existing metadata")
|
|
def step_create_state_with_metadata(context: Context):
|
|
"""Create GraphState with existing metadata."""
|
|
context.state_instances["with_metadata"] = GraphState(metadata={"key1": "value1", "nested": {"a": 1}})
|
|
|
|
|
|
@given("I have dictionary updates for merge mode")
|
|
def step_create_dict_updates_merge(context: Context):
|
|
"""Create dictionary updates for merge mode."""
|
|
context.test_data["dict_updates"] = {"metadata": {"key2": "value2", "nested": {"b": 2}}}
|
|
|
|
|
|
@when("I update the state with MERGE mode")
|
|
def step_update_state_merge_mode(context: Context):
|
|
"""Update state using MERGE mode."""
|
|
state = (
|
|
context.state_instances.get("with_metadata")
|
|
or context.state_instances.get("with_messages")
|
|
or context.state_instances.get("simple")
|
|
or context.state_instances.get("main")
|
|
)
|
|
updates = (
|
|
context.test_data.get("dict_updates")
|
|
or context.test_data.get("list_updates")
|
|
or context.test_data.get("simple_updates")
|
|
)
|
|
state.update(updates, StateUpdateMode.MERGE)
|
|
|
|
|
|
@then("dictionaries should be merged properly")
|
|
def step_verify_dict_merge(context: Context):
|
|
"""Verify dictionaries were merged properly."""
|
|
state = context.state_instances["with_metadata"]
|
|
assert state.metadata["key1"] == "value1" # Original value preserved
|
|
assert state.metadata["key2"] == "value2" # New value added
|
|
# Note: Python dict.update() replaces nested dicts completely
|
|
assert state.metadata["nested"]["b"] == 2 # New nested value added
|
|
assert "a" not in state.metadata["nested"] # Original nested value was replaced
|
|
|
|
|
|
@then("existing values should be preserved")
|
|
def step_verify_values_preserved(context: Context):
|
|
"""Verify existing values were preserved."""
|
|
state = context.state_instances["with_metadata"]
|
|
assert "key1" in state.metadata # Top-level keys are preserved
|
|
|
|
|
|
@given("I have a GraphState with existing messages")
|
|
def step_create_state_with_messages(context: Context):
|
|
"""Create GraphState with existing messages."""
|
|
context.state_instances["with_messages"] = GraphState(messages=[{"id": 1, "content": "first"}])
|
|
|
|
|
|
@given("I have list updates for merge mode")
|
|
def step_create_list_updates_merge(context: Context):
|
|
"""Create list updates for merge mode."""
|
|
context.test_data["list_updates"] = {"messages": [{"id": 2, "content": "second"}]}
|
|
|
|
|
|
@then("lists should be extended properly")
|
|
def step_verify_list_extend(context: Context):
|
|
"""Verify lists were extended properly."""
|
|
state = context.state_instances["with_messages"]
|
|
assert len(state.messages) == 2
|
|
assert state.messages[0]["id"] == 1
|
|
assert state.messages[1]["id"] == 2
|
|
|
|
|
|
@then("existing messages should be preserved")
|
|
def step_verify_messages_preserved(context: Context):
|
|
"""Verify existing messages were preserved."""
|
|
state = context.state_instances["with_messages"]
|
|
assert state.messages[0]["content"] == "first"
|
|
|
|
|
|
@given("I have a GraphState with simple values")
|
|
def step_create_state_simple_values(context: Context):
|
|
"""Create GraphState with simple values."""
|
|
context.state_instances["simple"] = GraphState(current_node="old_node", execution_count=5)
|
|
|
|
|
|
@given("I have simple value updates")
|
|
def step_create_simple_value_updates(context: Context):
|
|
"""Create simple value updates."""
|
|
context.test_data["simple_updates"] = {
|
|
"current_node": "new_node",
|
|
"execution_count": 10,
|
|
}
|
|
|
|
|
|
@then("simple values should be replaced")
|
|
def step_verify_simple_values_replaced(context: Context):
|
|
"""Verify simple values were replaced."""
|
|
state = context.state_instances["simple"]
|
|
assert state.current_node == "new_node"
|
|
assert state.execution_count == 10
|
|
|
|
|
|
@given("I have a GraphState with existing list data")
|
|
def step_create_state_with_list_data(context: Context):
|
|
"""Create GraphState with existing list data."""
|
|
context.state_instances["list_data"] = GraphState(messages=[{"id": 1}])
|
|
|
|
|
|
@given("I have new items to append")
|
|
def step_create_items_to_append(context: Context):
|
|
"""Create items to append."""
|
|
context.test_data["append_updates"] = {"messages": [{"id": 2}, {"id": 3}]}
|
|
context.test_data["single_append"] = {"messages": {"id": 4}}
|
|
|
|
|
|
@when("I update the state with APPEND mode")
|
|
def step_update_state_append_mode(context: Context):
|
|
"""Update state using APPEND mode."""
|
|
state = context.state_instances["list_data"]
|
|
# First append list items
|
|
state.update(context.test_data["append_updates"], StateUpdateMode.APPEND)
|
|
|
|
|
|
@then("items should be appended to lists")
|
|
def step_verify_items_appended(context: Context):
|
|
"""Verify items were appended to lists."""
|
|
state = context.state_instances["list_data"]
|
|
assert len(state.messages) == 3
|
|
assert state.messages[0]["id"] == 1
|
|
assert state.messages[1]["id"] == 2
|
|
assert state.messages[2]["id"] == 3
|
|
|
|
|
|
@then("single items should be appended as well")
|
|
def step_verify_single_items_appended(context: Context):
|
|
"""Verify single items can be appended."""
|
|
state = context.state_instances["list_data"]
|
|
# Append single item
|
|
state.update(context.test_data["single_append"], StateUpdateMode.APPEND)
|
|
assert len(state.messages) == 4
|
|
assert state.messages[3]["id"] == 4
|
|
|
|
|
|
@given("I have a GraphState with various data")
|
|
def step_create_state_with_various_data(context: Context):
|
|
"""Create GraphState with various data types."""
|
|
context.state_instances["various"] = GraphState(
|
|
messages=[{"role": "user"}],
|
|
metadata={"key": "value"},
|
|
current_node="test_node",
|
|
execution_count=42,
|
|
error="test_error",
|
|
)
|
|
|
|
|
|
@when("I convert the state to dictionary")
|
|
def step_convert_state_to_dict(context: Context):
|
|
"""Convert state to dictionary."""
|
|
state = context.state_instances["various"]
|
|
context.results["state_dict"] = state.to_dict()
|
|
|
|
|
|
@then("all fields should be present in the dictionary")
|
|
def step_verify_all_fields_present(context: Context):
|
|
"""Verify all fields are present in dictionary."""
|
|
state_dict = context.results["state_dict"]
|
|
required_fields = [
|
|
"messages",
|
|
"metadata",
|
|
"current_node",
|
|
"execution_count",
|
|
"error",
|
|
]
|
|
for field in required_fields:
|
|
assert field in state_dict
|
|
|
|
|
|
@then("the dictionary should match expected structure")
|
|
def step_verify_dict_structure(context: Context):
|
|
"""Verify dictionary structure matches expected."""
|
|
state_dict = context.results["state_dict"]
|
|
assert state_dict["messages"] == [{"role": "user"}]
|
|
assert state_dict["metadata"] == {"key": "value"}
|
|
assert state_dict["current_node"] == "test_node"
|
|
assert state_dict["execution_count"] == 42
|
|
assert state_dict["error"] == "test_error"
|
|
|
|
|
|
@given("I have a state dictionary")
|
|
def step_create_state_dictionary(context: Context):
|
|
"""Create a state dictionary."""
|
|
context.test_data["state_dict"] = {
|
|
"messages": [{"role": "assistant"}],
|
|
"metadata": {"source": "test"},
|
|
"current_node": "end_node",
|
|
"execution_count": 100,
|
|
"error": None,
|
|
}
|
|
|
|
|
|
@when("I create a GraphState from the dictionary")
|
|
def step_create_state_from_dict(context: Context):
|
|
"""Create GraphState from dictionary."""
|
|
state_dict = context.test_data["state_dict"]
|
|
context.state_instances["from_dict"] = GraphState.from_dict(state_dict)
|
|
|
|
|
|
@then("the GraphState should have correct field values")
|
|
def step_verify_state_from_dict(context: Context):
|
|
"""Verify GraphState created from dictionary has correct values."""
|
|
state = context.state_instances["from_dict"]
|
|
assert state.messages == [{"role": "assistant"}]
|
|
assert state.metadata == {"source": "test"}
|
|
assert state.current_node == "end_node"
|
|
assert state.execution_count == 100
|
|
assert state.error is None
|
|
|
|
|
|
@then("all data should be properly initialized")
|
|
def step_verify_data_initialized(context: Context):
|
|
"""Verify all data is properly initialized."""
|
|
state = context.state_instances["from_dict"]
|
|
assert isinstance(state.messages, list)
|
|
assert isinstance(state.metadata, dict)
|
|
assert isinstance(state.execution_count, int)
|
|
|
|
|
|
@given("I specify a checkpoint directory path")
|
|
def step_specify_checkpoint_directory(context: Context):
|
|
"""Specify a checkpoint directory path."""
|
|
context.test_data["checkpoint_dir"] = Path(tempfile.mkdtemp())
|
|
|
|
|
|
@when("I create a StateManager with checkpointing")
|
|
def step_create_state_manager_with_checkpointing(context: Context):
|
|
"""Create StateManager with checkpointing enabled."""
|
|
checkpoint_dir = context.test_data["checkpoint_dir"]
|
|
context.state_managers["with_checkpointing"] = StateManager(checkpoint_dir=checkpoint_dir)
|
|
|
|
|
|
@then("the checkpoint directory should be created")
|
|
def step_verify_checkpoint_dir_created(context: Context):
|
|
"""Verify checkpoint directory was created."""
|
|
checkpoint_dir = context.test_data["checkpoint_dir"]
|
|
assert checkpoint_dir.exists()
|
|
assert checkpoint_dir.is_dir()
|
|
|
|
|
|
@then("checkpointing should be enabled")
|
|
def step_verify_checkpointing_enabled(context: Context):
|
|
"""Verify checkpointing is enabled."""
|
|
manager = context.state_managers["with_checkpointing"]
|
|
assert manager.checkpoint_dir is not None
|
|
|
|
|
|
@given("I create a StateManager without time travel")
|
|
def step_create_state_manager_no_time_travel(context: Context):
|
|
"""Create StateManager without time travel."""
|
|
context.state_managers["no_time_travel"] = StateManager(enable_time_travel=False)
|
|
|
|
|
|
@when("I try to use time travel functionality")
|
|
def step_try_time_travel(context: Context):
|
|
"""Try to use time travel functionality."""
|
|
manager = context.state_managers["no_time_travel"]
|
|
context.results["time_travel_result"] = manager.time_travel(1)
|
|
|
|
|
|
@then("time travel should return None")
|
|
def step_verify_time_travel_none(context: Context):
|
|
"""Verify time travel returns None."""
|
|
assert context.results["time_travel_result"] is None
|
|
|
|
|
|
@then("no history should be maintained")
|
|
def step_verify_no_history(context: Context):
|
|
"""Verify no history is maintained."""
|
|
manager = context.state_managers["no_time_travel"]
|
|
assert not manager.enable_time_travel
|
|
assert len(manager.history) == 0
|
|
|
|
|
|
@given("I have a StateManager with time travel enabled")
|
|
def step_create_state_manager_time_travel(context: Context):
|
|
"""Create StateManager with time travel enabled."""
|
|
context.state_managers["time_travel"] = StateManager(enable_time_travel=True)
|
|
|
|
|
|
@given("I set a small max history size")
|
|
def step_set_small_history_size(context: Context):
|
|
"""Set a small max history size."""
|
|
manager = context.state_managers["time_travel"]
|
|
manager.max_history_size = 3
|
|
|
|
|
|
@when("I make many state updates")
|
|
def step_make_many_updates(context: Context):
|
|
"""Make many state updates."""
|
|
manager = context.state_managers["time_travel"]
|
|
for i in range(5):
|
|
manager.update_state({"execution_count": i}, node_id=f"node_{i}")
|
|
|
|
|
|
@then("the history should be trimmed to max size")
|
|
def step_verify_history_trimmed(context: Context):
|
|
"""Verify history is trimmed to max size."""
|
|
manager = context.state_managers["time_travel"]
|
|
assert len(manager.history) == 3
|
|
|
|
|
|
@then("older snapshots should be removed")
|
|
def step_verify_older_snapshots_removed(context: Context):
|
|
"""Verify older snapshots are removed."""
|
|
manager = context.state_managers["time_travel"]
|
|
# Should only have the last 3 snapshots
|
|
assert len(manager.history) <= manager.max_history_size
|
|
|
|
|
|
@given("I have a StateManager with checkpointing enabled")
|
|
def step_create_state_manager_checkpointing_enabled(context: Context):
|
|
"""Create StateManager with checkpointing enabled."""
|
|
checkpoint_dir = Path(tempfile.mkdtemp())
|
|
context.test_data["checkpoint_dir"] = checkpoint_dir
|
|
context.state_managers["checkpointing"] = StateManager(checkpoint_dir=checkpoint_dir)
|
|
|
|
|
|
@given("I set a small checkpoint interval")
|
|
def step_set_small_checkpoint_interval(context: Context):
|
|
"""Set a small checkpoint interval."""
|
|
manager = context.state_managers["checkpointing"]
|
|
manager.checkpoint_interval = 2
|
|
|
|
|
|
@when("I make multiple state updates")
|
|
def step_make_multiple_updates(context: Context):
|
|
"""Make multiple state updates."""
|
|
manager = context.state_managers["checkpointing"]
|
|
for i in range(3):
|
|
manager.update_state({"execution_count": i})
|
|
|
|
|
|
@then("checkpoints should be saved automatically")
|
|
def step_verify_checkpoints_saved(context: Context):
|
|
"""Verify checkpoints are saved automatically."""
|
|
checkpoint_dir = context.test_data["checkpoint_dir"]
|
|
checkpoint_files = list(checkpoint_dir.glob("checkpoint_*.json"))
|
|
assert len(checkpoint_files) > 0
|
|
|
|
|
|
@then("checkpoint files should be created")
|
|
def step_verify_checkpoint_files_created(context: Context):
|
|
"""Verify checkpoint files are created."""
|
|
checkpoint_dir = context.test_data["checkpoint_dir"]
|
|
checkpoint_files = list(checkpoint_dir.glob("checkpoint_*.json"))
|
|
assert len(checkpoint_files) >= 1
|
|
|
|
# Verify file content
|
|
with open(checkpoint_files[0], "r") as f:
|
|
checkpoint_data = json.load(f)
|
|
assert "state" in checkpoint_data
|
|
assert "timestamp" in checkpoint_data
|
|
assert "update_count" in checkpoint_data
|
|
|
|
|
|
@given("I have a StateManager with a checkpoint file")
|
|
def step_create_state_manager_with_checkpoint(context: Context):
|
|
"""Create StateManager with a checkpoint file."""
|
|
checkpoint_dir = Path(tempfile.mkdtemp())
|
|
context.test_data["checkpoint_dir"] = checkpoint_dir
|
|
|
|
# Create a checkpoint file
|
|
checkpoint_data = {
|
|
"state": {
|
|
"messages": [{"role": "system", "content": "loaded"}],
|
|
"metadata": {"loaded": True},
|
|
"current_node": "loaded_node",
|
|
"execution_count": 999,
|
|
"error": None,
|
|
},
|
|
"timestamp": "20240101_120000",
|
|
"update_count": 42,
|
|
}
|
|
|
|
checkpoint_file = checkpoint_dir / "checkpoint_test.json"
|
|
with open(checkpoint_file, "w") as f:
|
|
json.dump(checkpoint_data, f)
|
|
|
|
context.test_data["checkpoint_file"] = checkpoint_file
|
|
context.state_managers["with_checkpoint"] = StateManager(checkpoint_dir=checkpoint_dir)
|
|
|
|
|
|
@when("I load the checkpoint")
|
|
def step_load_checkpoint(context: Context):
|
|
"""Load the checkpoint."""
|
|
manager = context.state_managers["with_checkpoint"]
|
|
checkpoint_file = context.test_data["checkpoint_file"]
|
|
manager.load_checkpoint(checkpoint_file)
|
|
|
|
|
|
@then("the state should be restored from checkpoint")
|
|
def step_verify_state_restored(context: Context):
|
|
"""Verify state is restored from checkpoint."""
|
|
manager = context.state_managers["with_checkpoint"]
|
|
state = manager.get_state()
|
|
assert state.messages == [{"role": "system", "content": "loaded"}]
|
|
assert state.metadata == {"loaded": True}
|
|
assert state.current_node == "loaded_node"
|
|
assert state.execution_count == 999
|
|
|
|
|
|
@then("the update count should be restored")
|
|
def step_verify_update_count_restored(context: Context):
|
|
"""Verify update count is restored."""
|
|
manager = context.state_managers["with_checkpoint"]
|
|
assert manager.update_count == 42
|
|
|
|
|
|
@then("the state stream should emit the loaded state")
|
|
def step_verify_state_stream_emits(context: Context):
|
|
"""Verify state stream emits the loaded state."""
|
|
manager = context.state_managers["with_checkpoint"]
|
|
# The state stream should have the loaded state as current value
|
|
current_state = manager.state_stream.value if hasattr(manager.state_stream, "value") else manager.get_state()
|
|
assert current_state.current_node == "loaded_node"
|
|
|
|
|
|
@given("I have multiple checkpoint files")
|
|
def step_create_multiple_checkpoint_files(context: Context):
|
|
"""Create multiple checkpoint files."""
|
|
checkpoint_dir = Path(tempfile.mkdtemp())
|
|
context.test_data["checkpoint_dir"] = checkpoint_dir
|
|
|
|
# Create multiple checkpoint files with different timestamps
|
|
import time
|
|
|
|
for i, timestamp in enumerate(["20240101_100000", "20240101_110000", "20240101_120000"]):
|
|
checkpoint_file = checkpoint_dir / f"checkpoint_{timestamp}.json"
|
|
checkpoint_data = {
|
|
"state": {"execution_count": i},
|
|
"timestamp": timestamp,
|
|
"update_count": i,
|
|
}
|
|
with open(checkpoint_file, "w") as f:
|
|
json.dump(checkpoint_data, f)
|
|
# Ensure different modification times
|
|
time.sleep(0.01)
|
|
|
|
context.state_managers["multi_checkpoint"] = StateManager(checkpoint_dir=checkpoint_dir)
|
|
|
|
|
|
@when("I get the latest checkpoint")
|
|
def step_get_latest_checkpoint(context: Context):
|
|
"""Get the latest checkpoint."""
|
|
manager = context.state_managers["multi_checkpoint"]
|
|
context.results["latest_checkpoint"] = manager.get_latest_checkpoint()
|
|
|
|
|
|
@then("the most recent checkpoint should be returned")
|
|
def step_verify_most_recent_checkpoint(context: Context):
|
|
"""Verify the most recent checkpoint is returned."""
|
|
latest_checkpoint = context.results["latest_checkpoint"]
|
|
assert latest_checkpoint is not None
|
|
assert "checkpoint_" in str(latest_checkpoint)
|
|
|
|
|
|
@then("it should be based on file modification time")
|
|
def step_verify_based_on_modification_time(context: Context):
|
|
"""Verify selection is based on file modification time."""
|
|
latest_checkpoint = context.results["latest_checkpoint"]
|
|
# Should be the last created file
|
|
assert latest_checkpoint.name.endswith("120000.json")
|
|
|
|
|
|
@given("I have a checkpoint directory with no files")
|
|
def step_create_empty_checkpoint_dir(context: Context):
|
|
"""Create empty checkpoint directory."""
|
|
checkpoint_dir = Path(tempfile.mkdtemp())
|
|
context.test_data["empty_checkpoint_dir"] = checkpoint_dir
|
|
context.state_managers["empty_checkpoint"] = StateManager(checkpoint_dir=checkpoint_dir)
|
|
|
|
|
|
@when("I get the latest checkpoint from empty directory")
|
|
def step_get_latest_checkpoint_empty(context: Context):
|
|
"""Get latest checkpoint from empty directory."""
|
|
manager = context.state_managers["empty_checkpoint"]
|
|
context.results["empty_latest"] = manager.get_latest_checkpoint()
|
|
|
|
|
|
@then("None should be returned for latest checkpoint")
|
|
def step_verify_none_returned(context: Context):
|
|
"""Verify None is returned."""
|
|
assert context.results["empty_latest"] is None
|
|
|
|
|
|
@given("I have made several state updates")
|
|
def step_make_several_updates(context: Context):
|
|
"""Make several state updates."""
|
|
manager = context.state_managers["time_travel"]
|
|
for i in range(3):
|
|
manager.update_state({"execution_count": i}, node_id=f"step_{i}")
|
|
|
|
|
|
@when("I travel back in time")
|
|
def step_travel_back_in_time(context: Context):
|
|
"""Travel back in time."""
|
|
manager = context.state_managers["time_travel"]
|
|
context.results["time_travel_state"] = manager.time_travel(1)
|
|
|
|
|
|
@then("the state should revert to previous version")
|
|
def step_verify_state_reverted(context: Context):
|
|
"""Verify state reverted to previous version."""
|
|
reverted_state = context.results["time_travel_state"]
|
|
assert reverted_state is not None
|
|
# Should be the second-to-last state
|
|
assert reverted_state.execution_count == 1
|
|
|
|
|
|
@then("the state stream should emit the historical state")
|
|
def step_verify_stream_emits_historical(context: Context):
|
|
"""Verify state stream emits historical state."""
|
|
manager = context.state_managers["time_travel"]
|
|
current_state = manager.get_state()
|
|
assert current_state.execution_count == 1
|
|
|
|
|
|
@given("I have a StateManager with limited history")
|
|
def step_create_state_manager_limited_history(context: Context):
|
|
"""Create StateManager with limited history."""
|
|
context.state_managers["limited_history"] = StateManager(enable_time_travel=True)
|
|
# Add limited history
|
|
manager = context.state_managers["limited_history"]
|
|
manager.update_state({"execution_count": 1})
|
|
|
|
|
|
@when("I try to travel back more steps than available")
|
|
def step_try_travel_back_too_far(context: Context):
|
|
"""Try to travel back more steps than available."""
|
|
manager = context.state_managers["limited_history"]
|
|
context.results["limited_time_travel"] = manager.time_travel(10)
|
|
|
|
|
|
@then("it should travel back to the earliest available state")
|
|
def step_verify_earliest_state(context: Context):
|
|
"""Verify it travels back to earliest available state."""
|
|
result = context.results["limited_time_travel"]
|
|
assert result is not None # Should not crash
|
|
|
|
|
|
@then("time travel should not cause errors")
|
|
def step_verify_no_errors(context: Context):
|
|
"""Verify no errors are caused."""
|
|
# If we get here without exceptions, the test passes
|
|
assert True
|
|
|
|
|
|
@given("I have a StateManager")
|
|
def step_create_basic_state_manager(context: Context):
|
|
"""Create a basic StateManager."""
|
|
context.state_managers["basic"] = StateManager()
|
|
|
|
|
|
@when("I get the state observable")
|
|
def step_get_state_observable(context: Context):
|
|
"""Get the state observable."""
|
|
manager = context.state_managers["basic"]
|
|
context.results["observable"] = manager.get_state_observable()
|
|
|
|
|
|
@then("it should return an RxPy observable")
|
|
def step_verify_observable_returned(context: Context):
|
|
"""Verify RxPy observable is returned."""
|
|
observable = context.results["observable"]
|
|
# Check that it has observable-like methods
|
|
assert hasattr(observable, "subscribe")
|
|
|
|
|
|
@then("it should emit state changes")
|
|
def step_verify_observable_emits(context: Context):
|
|
"""Verify observable emits state changes."""
|
|
manager = context.state_managers["basic"]
|
|
observable = context.results["observable"]
|
|
|
|
# Subscribe to capture emissions
|
|
emissions = []
|
|
observable.subscribe(lambda state: emissions.append(state))
|
|
|
|
# Make a state update
|
|
manager.update_state({"execution_count": 999})
|
|
|
|
# Should have emitted the updated state
|
|
assert len(emissions) > 0
|
|
|
|
|
|
@given("I have a StateManager with some history")
|
|
def step_create_state_manager_with_history(context: Context):
|
|
"""Create StateManager with some history."""
|
|
context.state_managers["with_history"] = StateManager(enable_time_travel=True)
|
|
manager = context.state_managers["with_history"]
|
|
# Add some history
|
|
for i in range(3):
|
|
manager.update_state({"execution_count": i})
|
|
|
|
|
|
@when("I clear the history")
|
|
def step_clear_history(context: Context):
|
|
"""Clear the history."""
|
|
manager = context.state_managers["with_history"]
|
|
manager.clear_history()
|
|
|
|
|
|
@then("the history should be empty")
|
|
def step_verify_history_empty(context: Context):
|
|
"""Verify history is empty."""
|
|
manager = context.state_managers["with_history"]
|
|
assert len(manager.history) == 0
|
|
|
|
|
|
@then("time travel should not be possible")
|
|
def step_verify_time_travel_not_possible(context: Context):
|
|
"""Verify time travel is not possible."""
|
|
manager = context.state_managers["with_history"]
|
|
result = manager.time_travel(1)
|
|
assert result is None
|
|
|
|
|
|
@given("I have a StateManager with modified state")
|
|
def step_create_state_manager_modified(context: Context):
|
|
"""Create StateManager with modified state."""
|
|
context.state_managers["modified"] = StateManager(enable_time_travel=True)
|
|
manager = context.state_managers["modified"]
|
|
# Modify the state
|
|
manager.update_state({"execution_count": 100, "current_node": "modified"})
|
|
manager.update_state({"metadata": {"modified": True}})
|
|
|
|
|
|
@when("I reset the state manager")
|
|
def step_reset_state_manager(context: Context):
|
|
"""Reset the state manager."""
|
|
manager = context.state_managers["modified"]
|
|
manager.reset()
|
|
|
|
|
|
@then("the state should return to initial values")
|
|
def step_verify_state_reset_to_initial(context: Context):
|
|
"""Verify state returns to initial values."""
|
|
manager = context.state_managers["modified"]
|
|
state = manager.get_state()
|
|
assert state.execution_count == 0
|
|
assert state.current_node is None
|
|
assert state.metadata == {}
|
|
assert state.messages == []
|
|
assert state.error is None
|
|
|
|
|
|
@then("the update count should reset to zero")
|
|
def step_verify_update_count_reset(context: Context):
|
|
"""Verify update count resets to zero."""
|
|
manager = context.state_managers["modified"]
|
|
assert manager.update_count == 0
|
|
|
|
|
|
@then("the history should be cleared")
|
|
def step_verify_history_cleared(context: Context):
|
|
"""Verify history is cleared."""
|
|
manager = context.state_managers["modified"]
|
|
assert len(manager.history) == 0
|
|
|
|
|
|
@then("the state stream should emit the reset state")
|
|
def step_verify_stream_emits_reset(context: Context):
|
|
"""Verify state stream emits reset state."""
|
|
manager = context.state_managers["modified"]
|
|
current_state = manager.get_state()
|
|
assert current_state.execution_count == 0
|
|
|
|
|
|
@given("I have a custom initial state")
|
|
def step_create_custom_initial_state(context: Context):
|
|
"""Create a custom initial state."""
|
|
context.test_data["custom_initial"] = GraphState(
|
|
messages=[{"role": "system", "content": "custom"}],
|
|
metadata={"custom": True},
|
|
current_node="custom_start",
|
|
execution_count=50,
|
|
)
|
|
|
|
|
|
@when("I reset with the custom initial state")
|
|
def step_reset_with_custom_initial(context: Context):
|
|
"""Reset with custom initial state."""
|
|
manager = context.state_managers["basic"]
|
|
custom_initial = context.test_data["custom_initial"]
|
|
manager.reset(custom_initial)
|
|
|
|
|
|
@then("the state should match the custom initial state")
|
|
def step_verify_state_matches_custom(context: Context):
|
|
"""Verify state matches custom initial state."""
|
|
manager = context.state_managers["basic"]
|
|
state = manager.get_state()
|
|
assert state.messages == [{"role": "system", "content": "custom"}]
|
|
assert state.metadata == {"custom": True}
|
|
assert state.current_node == "custom_start"
|
|
assert state.execution_count == 50
|
|
|
|
|
|
@then("all counters should be reset")
|
|
def step_verify_all_counters_reset(context: Context):
|
|
"""Verify all counters are reset."""
|
|
manager = context.state_managers["basic"]
|
|
assert manager.update_count == 0
|
|
assert len(manager.history) == 0
|