Coverage for tests/graph/test_history.py: 96.61%

1# pyright: reportPrivateUsage=false 

2from __future__ import annotations as _annotations (empty)

3 

4import json (empty)

5from dataclasses import dataclass (empty)

6from datetime import datetime, timezone (empty)

7 

8import pytest (empty)

9from dirty_equals import IsStr (empty)

10from inline_snapshot import snapshot (empty)

11 

12from pydantic_graph import BaseNode, End, EndStep, Graph, GraphRunContext, GraphSetupError, NodeStep (empty)

13 

14from ..conftest import IsFloat, IsNow (empty)

15 

16pytestmark = pytest.mark.anyio (empty)

17 

18 

19@dataclass (empty)

20class MyState: (empty)

21 x: int (empty)

22 y: str (empty)

23 

24 

25@dataclass (empty)

26class Foo(BaseNode[MyState]): (empty)

27 async def run(self, ctx: GraphRunContext[MyState]) -> Bar: (empty)

28 ctx.state.x += 1 (empty)

29 return Bar() (empty)

30 

31 

32@dataclass (empty)

33class Bar(BaseNode[MyState, None, int]): (empty)

34 async def run(self, ctx: GraphRunContext[MyState]) -> End[int]: (empty)

35 ctx.state.y += 'y' (empty)

36 return End(ctx.state.x * 2) (empty)

37 

38 

39@pytest.mark.parametrize( (empty)

40 'graph', 

41 [ 

42 Graph(nodes=(Foo, Bar), state_type=MyState, run_end_type=int), 

43 Graph(nodes=(Foo, Bar), state_type=MyState), 

44 Graph(nodes=(Foo, Bar), run_end_type=int), 

45 Graph(nodes=(Foo, Bar)), 

46 ], 

47) 

48async def test_dump_load_history(graph: Graph[MyState, None, int]): (empty)

49 result, history = await graph.run(Foo(), state=MyState(1, '')) (empty)

50 assert result == snapshot(4) (empty)

51 assert history == snapshot( (empty)

52 [ 

53 NodeStep(state=MyState(x=2, y=''), node=Foo(), start_ts=IsNow(tz=timezone.utc), duration=IsFloat()), 

54 NodeStep(state=MyState(x=2, y='y'), node=Bar(), start_ts=IsNow(tz=timezone.utc), duration=IsFloat()), 

55 EndStep(result=End(4), ts=IsNow(tz=timezone.utc)), 

56 ] 

57 ) 

58 history_json = graph.dump_history(history) (empty)

59 assert json.loads(history_json) == snapshot( (empty)

60 [ 

61 { 

62 'state': {'x': 2, 'y': ''}, 

63 'node': {'node_id': 'Foo'}, 

64 'start_ts': IsStr(regex=r'20\d\d-\d\d-\d\dT.+'), 

65 'duration': IsFloat(), 

66 'kind': 'node', 

67 }, 

68 { 

69 'state': {'x': 2, 'y': 'y'}, 

70 'node': {'node_id': 'Bar'}, 

71 'start_ts': IsStr(regex=r'20\d\d-\d\d-\d\dT.+'), 

72 'duration': IsFloat(), 

73 'kind': 'node', 

74 }, 

75 {'result': {'data': 4}, 'ts': IsStr(regex=r'20\d\d-\d\d-\d\dT.+'), 'kind': 'end'}, 

76 ] 

77 ) 

78 history_loaded = graph.load_history(history_json) (empty)

79 assert history == history_loaded (empty)

80 

81 custom_history = [ (empty)

82 { 

83 'state': {'x': 2, 'y': ''}, 

84 'node': {'node_id': 'Foo'}, 

85 'start_ts': '2025-01-01T00:00:00Z', 

86 'duration': 123, 

87 'kind': 'node', 

88 }, 

89 {'result': {'data': '42'}, 'ts': '2025-01-01T00:00:00Z', 'kind': 'end'}, 

90 ] 

91 history_loaded = graph.load_history(json.dumps(custom_history)) (empty)

92 assert history_loaded == snapshot( (empty)

93 [ 

94 NodeStep( 

95 state=MyState(x=2, y=''), 

96 node=Foo(), 

97 start_ts=datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc), 

98 duration=123.0, 

99 ), 

100 EndStep(result=End(data=42), ts=datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc)), 

101 ] 

102 ) 

103 

104 

105def test_one_node(): (empty)

106 @dataclass (empty)

107 class MyNode(BaseNode[None, None, int]): (empty)

108 async def run(self, ctx: GraphRunContext) -> End[int]: (empty)

109 return End(123) 

110 

111 g = Graph(nodes=[MyNode]) (empty)

112 

113 custom_history = [ (empty)

114 {'result': {'data': '123'}, 'ts': '2025-01-01T00:00:00Z', 'kind': 'end'}, 

115 ] 

116 history_loaded = g.load_history(json.dumps(custom_history)) (empty)

117 assert history_loaded == snapshot( (empty)

118 [ 

119 EndStep(result=End(data=123), ts=datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc)), 

120 ] 

121 ) 

122 

123 

124def test_no_generic_arg(): (empty)

125 @dataclass (empty)

126 class NoGenericArgsNode(BaseNode): (empty)

127 async def run(self, ctx: GraphRunContext) -> NoGenericArgsNode: (empty)

128 return NoGenericArgsNode() 

129 

130 g = Graph(nodes=[NoGenericArgsNode]) (empty)

131 assert g._get_state_type() is type(None) (empty)

132 with pytest.raises(GraphSetupError, match='Could not infer run end type from nodes, please set `run_end_type`.'): (empty)

133 g._get_run_end_type() (empty)

134 

135 g = Graph(nodes=[NoGenericArgsNode], run_end_type=None) # pyright: ignore[reportArgumentType] (empty)

136 assert g._get_run_end_type() is None (empty)

137 

138 custom_history = [ (empty)

139 {'result': {'data': None}, 'ts': '2025-01-01T00:00:00Z', 'kind': 'end'}, 

140 ] 

141 history_loaded = g.load_history(json.dumps(custom_history)) (empty)

142 assert history_loaded == snapshot( (empty)

143 [ 

144 EndStep(result=End(data=None), ts=datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc)), 

145 ] 

146 )