Coverage for tests/graph/test_graph.py: 99.54%

1# pyright: reportPrivateUsage=false 

2from __future__ import annotations as _annotations (empty)

3 

4from dataclasses import dataclass (empty)

5from datetime import timezone (empty)

6from functools import cache (empty)

7from typing import Union (empty)

8 

9import pytest (empty)

10from dirty_equals import IsStr (empty)

11from inline_snapshot import snapshot (empty)

12 

13from pydantic_graph import ( (empty)

14 BaseNode, 

15 End, 

16 EndSnapshot, 

17 FullStatePersistence, 

18 Graph, 

19 GraphRunContext, 

20 GraphRuntimeError, 

21 GraphSetupError, 

22 NodeSnapshot, 

23 SimpleStatePersistence, 

24) 

25 

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

27 

28pytestmark = pytest.mark.anyio (empty)

29 

30 

31@dataclass (empty)

32class Float2String(BaseNode): (empty)

33 input_data: float (empty)

34 

35 async def run(self, ctx: GraphRunContext) -> String2Length: (empty)

36 return String2Length(str(self.input_data)) (empty)

37 

38 

39@dataclass (empty)

40class String2Length(BaseNode): (empty)

41 input_data: str (empty)

42 

43 async def run(self, ctx: GraphRunContext) -> Double: (empty)

44 return Double(len(self.input_data)) (empty)

45 

46 

47@dataclass (empty)

48class Double(BaseNode[None, None, int]): (empty)

49 input_data: int (empty)

50 

51 async def run(self, ctx: GraphRunContext) -> Union[String2Length, End[int]]: # noqa: UP007 (empty)

52 if self.input_data == 7: (empty)

53 return String2Length('x' * 21) (empty)

54 else: 

55 return End(self.input_data * 2) (empty)

56 

57 

58async def test_graph(): (empty)

59 my_graph = Graph(nodes=(Float2String, String2Length, Double)) (empty)

60 assert my_graph.name is None (empty)

61 assert my_graph.inferred_types == (type(None), int) (empty)

62 result = await my_graph.run(Float2String(3.14)) (empty)

63 # len('3.14') * 2 == 8 

64 assert result.output == 8 (empty)

65 assert my_graph.name == 'my_graph' (empty)

66 

67 

68async def test_graph_history(mock_snapshot_id: object): (empty)

69 my_graph = Graph[None, None, int](nodes=(Float2String, String2Length, Double)) (empty)

70 assert my_graph.name is None (empty)

71 assert my_graph.inferred_types == (type(None), int) (empty)

72 sp = FullStatePersistence() (empty)

73 result = await my_graph.run(Float2String(3.14), persistence=sp) (empty)

74 # len('3.14') * 2 == 8 

75 assert result.output == 8 (empty)

76 assert my_graph.name == 'my_graph' (empty)

77 assert sp.history == snapshot( (empty)

78 [ 

79 NodeSnapshot( 

80 state=None, 

81 node=Float2String(input_data=3.14), 

82 start_ts=IsNow(tz=timezone.utc), 

83 status='success', 

84 id='Float2String:1', 

85 duration=IsFloat(), 

86 ), 

87 NodeSnapshot( 

88 state=None, 

89 node=String2Length(input_data='3.14'), 

90 start_ts=IsNow(tz=timezone.utc), 

91 status='success', 

92 id='String2Length:2', 

93 duration=IsFloat(), 

94 ), 

95 NodeSnapshot( 

96 state=None, 

97 node=Double(input_data=4), 

98 start_ts=IsNow(tz=timezone.utc), 

99 status='success', 

100 id='Double:3', 

101 duration=IsFloat(), 

102 ), 

103 EndSnapshot(state=None, result=End(data=8), ts=IsNow(tz=timezone.utc), id='end:4'), 

104 ] 

105 ) 

106 sp = FullStatePersistence() (empty)

107 result = await my_graph.run(Float2String(3.14159), persistence=sp) (empty)

108 # len('3.14159') == 7, 21 * 2 == 42 

109 assert result.output == 42 (empty)

110 assert sp.history == snapshot( (empty)

111 [ 

112 NodeSnapshot( 

113 state=None, 

114 node=Float2String(input_data=3.14159), 

115 start_ts=IsNow(tz=timezone.utc), 

116 status='success', 

117 id='Float2String:5', 

118 duration=IsFloat(), 

119 ), 

120 NodeSnapshot( 

121 state=None, 

122 node=String2Length(input_data='3.14159'), 

123 start_ts=IsNow(tz=timezone.utc), 

124 status='success', 

125 id='String2Length:6', 

126 duration=IsFloat(), 

127 ), 

128 NodeSnapshot( 

129 state=None, 

130 node=Double(input_data=7), 

131 start_ts=IsNow(tz=timezone.utc), 

132 status='success', 

133 id='Double:7', 

134 duration=IsFloat(), 

135 ), 

136 NodeSnapshot( 

137 state=None, 

138 node=String2Length(input_data='xxxxxxxxxxxxxxxxxxxxx'), 

139 start_ts=IsNow(tz=timezone.utc), 

140 status='success', 

141 id='String2Length:8', 

142 duration=IsFloat(), 

143 ), 

144 NodeSnapshot( 

145 state=None, 

146 node=Double(input_data=21), 

147 start_ts=IsNow(tz=timezone.utc), 

148 status='success', 

149 id='Double:9', 

150 duration=IsFloat(), 

151 ), 

152 EndSnapshot(state=None, result=End(data=42), ts=IsNow(tz=timezone.utc), id='end:10'), 

153 ] 

154 ) 

155 assert [e.node for e in sp.history] == snapshot( (empty)

156 [ 

157 Float2String(input_data=3.14159), 

158 String2Length(input_data='3.14159'), 

159 Double(input_data=7), 

160 String2Length(input_data='xxxxxxxxxxxxxxxxxxxxx'), 

161 Double(input_data=21), 

162 End(data=42), 

163 ] 

164 ) 

165 

166 

167def test_one_bad_node(): (empty)

168 class Float2String(BaseNode): (empty)

169 async def run(self, ctx: GraphRunContext) -> String2Length: (empty)

170 raise NotImplementedError() 

171 

172 class String2Length(BaseNode[None, None, None]): # pyright: ignore[reportUnusedClass] (empty)

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

174 raise NotImplementedError() 

175 

176 with pytest.raises(GraphSetupError) as exc_info: (empty)

177 Graph(nodes=(Float2String,)) (empty)

178 

179 assert exc_info.value.message == snapshot( (empty)

180 '`String2Length` is referenced by `Float2String` but not included in the graph.' 

181 ) 

182 

183 

184def test_two_bad_nodes(): (empty)

185 class Foo(BaseNode): (empty)

186 input_data: float (empty)

187 

188 async def run(self, ctx: GraphRunContext) -> Union[Bar, Spam]: # noqa: UP007 (empty)

189 raise NotImplementedError() 

190 

191 class Bar(BaseNode[None, None, None]): (empty)

192 input_data: str (empty)

193 

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

195 raise NotImplementedError() 

196 

197 class Spam(BaseNode[None, None, None]): (empty)

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

199 raise NotImplementedError() 

200 

201 with pytest.raises(GraphSetupError) as exc_info: (empty)

202 Graph(nodes=(Foo,)) (empty)

203 

204 assert exc_info.value.message == snapshot("""\ (empty)

205Nodes are referenced in the graph but not included in the graph: 

206 `Bar` is referenced by `Foo` 

207 `Spam` is referenced by `Foo`\ 

208""") 

209 

210 

211def test_three_bad_nodes_separate(): (empty)

212 class Foo(BaseNode): (empty)

213 input_data: float (empty)

214 

215 async def run(self, ctx: GraphRunContext) -> Eggs: (empty)

216 raise NotImplementedError() 

217 

218 class Bar(BaseNode[None, None, None]): (empty)

219 input_data: str (empty)

220 

221 async def run(self, ctx: GraphRunContext) -> Eggs: (empty)

222 raise NotImplementedError() 

223 

224 class Spam(BaseNode[None, None, None]): (empty)

225 async def run(self, ctx: GraphRunContext) -> Eggs: (empty)

226 raise NotImplementedError() 

227 

228 class Eggs(BaseNode[None, None, None]): (empty)

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

230 raise NotImplementedError() 

231 

232 with pytest.raises(GraphSetupError) as exc_info: (empty)

233 Graph(nodes=(Foo, Bar, Spam)) (empty)

234 

235 assert exc_info.value.message == snapshot( (empty)

236 '`Eggs` is referenced by `Foo`, `Bar`, and `Spam` but not included in the graph.' 

237 ) 

238 

239 

240def test_duplicate_id(): (empty)

241 class Foo(BaseNode): (empty)

242 async def run(self, ctx: GraphRunContext) -> Bar: (empty)

243 raise NotImplementedError() 

244 

245 class Bar(BaseNode[None, None, None]): (empty)

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

247 raise NotImplementedError() 

248 

249 @classmethod (empty)

250 @cache (empty)

251 def get_node_id(cls) -> str: (empty)

252 return 'Foo' (empty)

253 

254 with pytest.raises(GraphSetupError) as exc_info: (empty)

255 Graph(nodes=(Foo, Bar)) (empty)

256 

257 assert exc_info.value.message == snapshot(IsStr(regex='Node ID `Foo` is not unique — found on <class.+')) (empty)

258 

259 

260async def test_run_node_not_in_graph(): (empty)

261 @dataclass (empty)

262 class Foo(BaseNode): (empty)

263 async def run(self, ctx: GraphRunContext) -> Bar: (empty)

264 return Bar() (empty)

265 

266 @dataclass (empty)

267 class Bar(BaseNode[None, None, None]): (empty)

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

269 return Spam() # type: ignore (empty)

270 

271 @dataclass (empty)

272 class Spam(BaseNode[None, None, None]): (empty)

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

274 raise NotImplementedError() 

275 

276 g = Graph(nodes=(Foo, Bar)) (empty)

277 with pytest.raises(GraphRuntimeError) as exc_info: (empty)

278 await g.run(Foo()) (empty)

279 

280 assert exc_info.value.message == snapshot('Node `test_run_node_not_in_graph.<locals>.Spam()` is not in the graph.') (empty)

281 

282 

283async def test_run_return_other(mock_snapshot_id: object): (empty)

284 @dataclass (empty)

285 class Foo(BaseNode): (empty)

286 async def run(self, ctx: GraphRunContext) -> Bar: (empty)

287 return Bar() (empty)

288 

289 @dataclass (empty)

290 class Bar(BaseNode[None, None, None]): (empty)

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

292 return 42 # type: ignore (empty)

293 

294 g = Graph(nodes=(Foo, Bar)) (empty)

295 assert g.inferred_types == (type(None), type(None)) (empty)

296 with pytest.raises(GraphRuntimeError) as exc_info: (empty)

297 await g.run(Foo()) (empty)

298 

299 assert exc_info.value.message == snapshot('Invalid node return type: `int`. Expected `BaseNode` or `End`.') (empty)

300 

301 

302async def test_iter(): (empty)

303 my_graph = Graph(nodes=(Float2String, String2Length, Double)) (empty)

304 assert my_graph.name is None (empty)

305 assert my_graph.inferred_types == (type(None), int) (empty)

306 node_reprs: list[str] = [] (empty)

307 async with my_graph.iter(Float2String(3.14)) as graph_iter: (empty)

308 assert repr(graph_iter) == snapshot('<GraphRun graph=my_graph>') (empty)

309 async for node in graph_iter: (empty)

310 node_reprs.append(repr(node)) (empty)

311 # len('3.14') * 2 == 8 

312 assert graph_iter.result (empty)

313 assert graph_iter.result.output == 8 (empty)

314 

315 assert node_reprs == snapshot(["String2Length(input_data='3.14')", 'Double(input_data=4)', 'End(data=8)']) (empty)

316 

317 

318async def test_iter_next(mock_snapshot_id: object): (empty)

319 @dataclass (empty)

320 class Foo(BaseNode): (empty)

321 async def run(self, ctx: GraphRunContext) -> Bar: (empty)

322 return Bar() (empty)

323 

324 @dataclass (empty)

325 class Bar(BaseNode): (empty)

326 async def run(self, ctx: GraphRunContext) -> Foo: (empty)

327 return Foo() (empty)

328 

329 g = Graph(nodes=(Foo, Bar)) (empty)

330 assert g.name is None (empty)

331 sp = FullStatePersistence() (empty)

332 async with g.iter(Foo(), persistence=sp) as run: (empty)

333 assert g.name == 'g' (empty)

334 n = await run.next() (empty)

335 assert n == Bar() (empty)

336 assert sp.history == snapshot( (empty)

337 [ 

338 NodeSnapshot( 

339 state=None, 

340 node=Foo(), 

341 start_ts=IsNow(tz=timezone.utc), 

342 duration=IsFloat(), 

343 status='success', 

344 id='Foo:1', 

345 ), 

346 NodeSnapshot(state=None, node=Bar(), id='Bar:2'), 

347 ] 

348 ) 

349 

350 assert isinstance(n, Bar) (empty)

351 n2 = await run.next() (empty)

352 assert n2 == Foo() (empty)

353 

354 assert sp.history == snapshot( (empty)

355 [ 

356 NodeSnapshot( 

357 state=None, 

358 node=Foo(), 

359 start_ts=IsNow(tz=timezone.utc), 

360 duration=IsFloat(), 

361 status='success', 

362 id='Foo:1', 

363 ), 

364 NodeSnapshot( 

365 state=None, 

366 node=Bar(), 

367 start_ts=IsNow(tz=timezone.utc), 

368 duration=IsFloat(), 

369 status='success', 

370 id='Bar:2', 

371 ), 

372 NodeSnapshot(state=None, node=Foo(), id='Foo:3'), 

373 ] 

374 ) 

375 

376 

377async def test_iter_next_error(mock_snapshot_id: object): (empty)

378 @dataclass (empty)

379 class Foo(BaseNode): (empty)

380 async def run(self, ctx: GraphRunContext) -> Bar: (empty)

381 return Bar() (empty)

382 

383 @dataclass (empty)

384 class Bar(BaseNode[None, None, None]): (empty)

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

386 return End(None) (empty)

387 

388 g = Graph(nodes=(Foo, Bar)) (empty)

389 sp = SimpleStatePersistence() (empty)

390 async with g.iter(Foo(), persistence=sp) as run: (empty)

391 n = await run.next() (empty)

392 assert n == snapshot(Bar()) (empty)

393 

394 assert isinstance(n, BaseNode) (empty)

395 n = await run.next() (empty)

396 assert n == snapshot(End(data=None)) (empty)

397 

398 with pytest.raises(TypeError, match=r'`next` must be called with a `BaseNode` instance, got End\(data=None\).'): (empty)

399 await run.next() (empty)

400 

401 

402async def test_next(mock_snapshot_id: object): (empty)

403 @dataclass (empty)

404 class Foo(BaseNode): (empty)

405 async def run(self, ctx: GraphRunContext) -> Bar: (empty)

406 return Bar() (empty)

407 

408 @dataclass (empty)

409 class Bar(BaseNode): (empty)

410 async def run(self, ctx: GraphRunContext) -> Foo: (empty)

411 return Foo() 

412 

413 g = Graph(nodes=(Foo, Bar)) (empty)

414 assert g.name is None (empty)

415 sp = FullStatePersistence() (empty)

416 with pytest.warns(DeprecationWarning, match='`next` is deprecated, use `async with graph.iter(...)'): (empty)

417 n = await g.next(Foo(), persistence=sp) # pyright: ignore[reportDeprecated] (empty)

418 assert n == Bar() (empty)

419 assert g.name == 'g' (empty)

420 assert sp.history == snapshot( (empty)

421 [ 

422 NodeSnapshot( 

423 state=None, 

424 node=Foo(), 

425 start_ts=IsNow(tz=timezone.utc), 

426 duration=IsFloat(), 

427 status='success', 

428 id='Foo:1', 

429 ), 

430 NodeSnapshot(state=None, node=Bar(), id='Bar:2'), 

431 ] 

432 ) 

433 

434 

435async def test_deps(mock_snapshot_id: object): (empty)

436 @dataclass (empty)

437 class Deps: (empty)

438 a: int (empty)

439 b: int (empty)

440 

441 @dataclass (empty)

442 class Foo(BaseNode[None, Deps]): (empty)

443 async def run(self, ctx: GraphRunContext[None, Deps]) -> Bar: (empty)

444 assert isinstance(ctx.deps, Deps) (empty)

445 return Bar() (empty)

446 

447 @dataclass (empty)

448 class Bar(BaseNode[None, Deps, int]): (empty)

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

450 assert isinstance(ctx.deps, Deps) (empty)

451 return End(123) (empty)

452 

453 g = Graph(nodes=(Foo, Bar)) (empty)

454 sp = FullStatePersistence() (empty)

455 result = await g.run(Foo(), deps=Deps(1, 2), persistence=sp) (empty)

456 

457 assert result.output == 123 (empty)

458 assert sp.history == snapshot( (empty)

459 [ 

460 NodeSnapshot( 

461 state=None, 

462 node=Foo(), 

463 start_ts=IsNow(tz=timezone.utc), 

464 duration=IsFloat(), 

465 status='success', 

466 id='Foo:1', 

467 ), 

468 NodeSnapshot( 

469 state=None, 

470 node=Bar(), 

471 start_ts=IsNow(tz=timezone.utc), 

472 duration=IsFloat(), 

473 status='success', 

474 id='Bar:2', 

475 ), 

476 EndSnapshot(state=None, result=End(data=123), ts=IsNow(tz=timezone.utc), id='end:3'), 

477 ] 

478 )