Spaces:
Sleeping
Sleeping
File size: 5,234 Bytes
f871fed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
"""
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"] == ""
@pytest.mark.asyncio
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"])
|