Spaces:
Sleeping
Sleeping
| """ | |
| Unit tests for the open_notebook.graphs module. | |
| This test suite focuses on testing graph structures, tools, and validation | |
| without heavy mocking of the actual processing logic. | |
| """ | |
| from datetime import datetime | |
| import pytest | |
| from open_notebook.graphs.prompt import PatternChainState, graph | |
| from open_notebook.graphs.tools import get_current_timestamp | |
| from open_notebook.graphs.transformation import ( | |
| TransformationState, | |
| run_transformation, | |
| ) | |
| from open_notebook.graphs.transformation import ( | |
| graph as transformation_graph, | |
| ) | |
| # ============================================================================ | |
| # TEST SUITE 1: Graph Tools | |
| # ============================================================================ | |
| class TestGraphTools: | |
| """Test suite for graph tool definitions.""" | |
| def test_get_current_timestamp_format(self): | |
| """Test timestamp tool returns correct format.""" | |
| timestamp = get_current_timestamp.func() | |
| assert isinstance(timestamp, str) | |
| assert len(timestamp) == 14 # YYYYMMDDHHmmss format | |
| assert timestamp.isdigit() | |
| def test_get_current_timestamp_validity(self): | |
| """Test timestamp represents valid datetime.""" | |
| timestamp = get_current_timestamp.func() | |
| # Parse it back to datetime to verify validity | |
| year = int(timestamp[0:4]) | |
| month = int(timestamp[4:6]) | |
| day = int(timestamp[6:8]) | |
| hour = int(timestamp[8:10]) | |
| minute = int(timestamp[10:12]) | |
| second = int(timestamp[12:14]) | |
| # Should be valid date components | |
| assert 2020 <= year <= 2100 | |
| assert 1 <= month <= 12 | |
| assert 1 <= day <= 31 | |
| assert 0 <= hour <= 23 | |
| assert 0 <= minute <= 59 | |
| assert 0 <= second <= 59 | |
| # Should parse as datetime | |
| dt = datetime.strptime(timestamp, "%Y%m%d%H%M%S") | |
| assert isinstance(dt, datetime) | |
| def test_get_current_timestamp_is_tool(self): | |
| """Test that function is properly decorated as a tool.""" | |
| # Check it has tool attributes | |
| assert hasattr(get_current_timestamp, "name") | |
| assert hasattr(get_current_timestamp, "description") | |
| # ============================================================================ | |
| # TEST SUITE 2: Prompt Graph State | |
| # ============================================================================ | |
| class TestPromptGraph: | |
| """Test suite for prompt pattern chain graph.""" | |
| def test_pattern_chain_state_structure(self): | |
| """Test PatternChainState structure and fields.""" | |
| state = PatternChainState( | |
| prompt="Test prompt", | |
| parser=None, | |
| input_text="Test input", | |
| output="" | |
| ) | |
| assert state["prompt"] == "Test prompt" | |
| assert state["parser"] is None | |
| assert state["input_text"] == "Test input" | |
| assert state["output"] == "" | |
| def test_prompt_graph_compilation(self): | |
| """Test that prompt graph compiles correctly.""" | |
| assert graph is not None | |
| # Graph should have the expected structure | |
| assert hasattr(graph, "invoke") | |
| assert hasattr(graph, "ainvoke") | |
| # ============================================================================ | |
| # TEST SUITE 3: Transformation Graph | |
| # ============================================================================ | |
| class TestTransformationGraph: | |
| """Test suite for transformation graph workflows.""" | |
| def test_transformation_state_structure(self): | |
| """Test TransformationState structure and fields.""" | |
| from unittest.mock import MagicMock | |
| from open_notebook.domain.notebook import Source | |
| from open_notebook.domain.transformation import Transformation | |
| mock_source = MagicMock(spec=Source) | |
| mock_transformation = MagicMock(spec=Transformation) | |
| state = TransformationState( | |
| input_text="Test text", | |
| source=mock_source, | |
| transformation=mock_transformation, | |
| output="" | |
| ) | |
| assert state["input_text"] == "Test text" | |
| assert state["source"] == mock_source | |
| assert state["transformation"] == mock_transformation | |
| assert state["output"] == "" | |
| async def test_run_transformation_assertion_no_content(self): | |
| """Test transformation raises assertion with no content.""" | |
| from unittest.mock import MagicMock | |
| from open_notebook.domain.transformation import Transformation | |
| mock_transformation = MagicMock(spec=Transformation) | |
| state = { | |
| "input_text": None, | |
| "transformation": mock_transformation, | |
| "source": None | |
| } | |
| config = {"configurable": {"model_id": None}} | |
| with pytest.raises(AssertionError, match="No content to transform"): | |
| await run_transformation(state, config) | |
| def test_transformation_graph_compilation(self): | |
| """Test that transformation graph compiles correctly.""" | |
| assert transformation_graph is not None | |
| assert hasattr(transformation_graph, "invoke") | |
| assert hasattr(transformation_graph, "ainvoke") | |
| if __name__ == "__main__": | |
| pytest.main([__file__, "-v"]) | |