1use crate::{
2 db::api::view::{filter_ops::Filter, StaticGraphViewOps},
3 prelude::{EdgeViewOps, Graph, GraphViewOps, NodeViewOps},
4};
5use std::ops::Range;
6
7#[cfg(feature = "search")]
8pub use crate::db::api::view::SearchableGraphOps;
9#[cfg(feature = "search")]
10use crate::prelude::IndexMutationOps;
11use crate::{
12 db::{
13 api::view::filter_ops::{EdgeSelect, NodeSelect},
14 graph::views::{
15 filter::{model::TryAsCompositeFilter, CreateFilter},
16 window_graph::WindowedGraph,
17 },
18 },
19 errors::GraphError,
20 prelude::TimeOps,
21};
22use raphtory_api::core::Direction;
23#[cfg(feature = "storage")]
24use {
25 crate::db::api::storage::graph::storage_ops::disk_storage::IntoGraph,
26 raphtory_storage::disk::DiskGraphStorage, tempfile::TempDir,
27};
28
29pub enum TestGraphVariants {
30 Graph,
31 PersistentGraph,
32 EventDiskGraph,
33 PersistentDiskGraph,
34}
35
36impl Into<Vec<TestGraphVariants>> for TestGraphVariants {
37 fn into(self) -> Vec<TestGraphVariants> {
38 vec![self]
39 }
40}
41
42pub enum TestVariants {
43 All,
44 EventOnly,
45 PersistentOnly,
46 NonDiskOnly,
47 DiskOnly,
48}
49
50impl From<TestVariants> for Vec<TestGraphVariants> {
51 fn from(variants: TestVariants) -> Self {
52 use TestGraphVariants::*;
53 match variants {
54 TestVariants::All => {
55 vec![Graph, PersistentGraph, EventDiskGraph, PersistentDiskGraph]
56 }
57 TestVariants::EventOnly => vec![Graph, EventDiskGraph],
58 TestVariants::PersistentOnly => vec![PersistentGraph, PersistentDiskGraph],
59 TestVariants::NonDiskOnly => vec![Graph, PersistentGraph],
60 TestVariants::DiskOnly => vec![EventDiskGraph, PersistentDiskGraph],
61 }
62 }
63}
64
65pub trait GraphTransformer {
66 type Return<G: StaticGraphViewOps>: StaticGraphViewOps;
67 fn apply<G: StaticGraphViewOps>(&self, graph: G) -> Self::Return<G>;
68}
69
70pub struct WindowGraphTransformer(pub Range<i64>);
71
72impl GraphTransformer for WindowGraphTransformer {
73 type Return<G: StaticGraphViewOps> = WindowedGraph<G>;
74 fn apply<G: StaticGraphViewOps>(&self, graph: G) -> Self::Return<G> {
75 graph.window(self.0.start, self.0.end)
76 }
77}
78
79pub trait ApplyFilter {
80 fn apply<G: StaticGraphViewOps>(&self, graph: G) -> Vec<String>;
81}
82
83pub struct FilterNodes<F: TryAsCompositeFilter + CreateFilter + Clone>(F);
84
85impl<F: TryAsCompositeFilter + CreateFilter + Clone> ApplyFilter for FilterNodes<F> {
86 fn apply<G: StaticGraphViewOps>(&self, graph: G) -> Vec<String> {
87 let mut results = graph
88 .filter(self.0.clone())
89 .unwrap()
90 .nodes()
91 .iter()
92 .map(|n| n.name())
93 .collect::<Vec<_>>();
94 results.sort();
95 results
96 }
97}
98
99pub struct SelectNodes<F: TryAsCompositeFilter + CreateFilter + Clone>(F);
100
101impl<F: TryAsCompositeFilter + CreateFilter + Clone> ApplyFilter for SelectNodes<F> {
102 fn apply<G: StaticGraphViewOps>(&self, graph: G) -> Vec<String> {
103 let mut results = graph
104 .nodes()
105 .select(self.0.clone())
106 .unwrap()
107 .iter()
108 .map(|n| n.name())
109 .collect::<Vec<_>>();
110 results.sort();
111 results
112 }
113}
114
115pub struct FilterNeighbours<F: TryAsCompositeFilter + CreateFilter + Clone>(F, String, Direction);
116
117impl<F: TryAsCompositeFilter + CreateFilter + Clone> ApplyFilter for FilterNeighbours<F> {
118 fn apply<G: StaticGraphViewOps>(&self, graph: G) -> Vec<String> {
119 let filter_applied = graph
120 .node(self.1.clone())
121 .unwrap()
122 .filter(self.0.clone())
123 .unwrap();
124
125 let mut results = match self.2 {
126 Direction::OUT => filter_applied.out_neighbours(),
127 Direction::IN => filter_applied.in_neighbours(),
128 Direction::BOTH => filter_applied.neighbours(),
129 }
130 .iter()
131 .map(|n| n.name())
132 .collect::<Vec<_>>();
133 results.sort();
134 results
135 }
136}
137
138pub struct SearchNodes<F: TryAsCompositeFilter + CreateFilter + Clone>(F);
139
140impl<F: TryAsCompositeFilter + CreateFilter + Clone> ApplyFilter for SearchNodes<F> {
141 fn apply<G: StaticGraphViewOps>(&self, graph: G) -> Vec<String> {
142 #[cfg(feature = "search")]
143 {
144 let mut results = graph
145 .search_nodes(self.0.clone(), 20, 0)
146 .unwrap()
147 .into_iter()
148 .map(|nv| nv.name())
149 .collect::<Vec<_>>();
150 results.sort();
151 return results;
152 }
153 #[cfg(not(feature = "search"))]
154 Vec::<String>::new()
155 }
156}
157
158pub struct FilterEdges<F: TryAsCompositeFilter + CreateFilter + Clone>(F);
159
160impl<F: TryAsCompositeFilter + CreateFilter + Clone> ApplyFilter for FilterEdges<F> {
161 fn apply<G: StaticGraphViewOps>(&self, graph: G) -> Vec<String> {
162 let mut results = graph
163 .filter(self.0.clone())
164 .unwrap()
165 .edges()
166 .iter()
167 .map(|e| format!("{}->{}", e.src().name(), e.dst().name()))
168 .collect::<Vec<_>>();
169 results.sort();
170 results
171 }
172}
173
174pub struct SelectEdges<F: TryAsCompositeFilter + CreateFilter + Clone>(F);
175
176impl<F: TryAsCompositeFilter + CreateFilter + Clone> ApplyFilter for SelectEdges<F> {
177 fn apply<G: StaticGraphViewOps>(&self, graph: G) -> Vec<String> {
178 let mut results = graph
179 .edges()
180 .select(self.0.clone())
181 .unwrap()
182 .into_iter()
183 .map(|e| format!("{}->{}", e.src().name(), e.dst().name()))
184 .collect::<Vec<_>>();
185 results.sort();
186 results
187 }
188}
189
190pub struct SearchEdges<F: TryAsCompositeFilter + CreateFilter + Clone>(F);
191
192impl<F: TryAsCompositeFilter + CreateFilter + Clone> ApplyFilter for SearchEdges<F> {
193 fn apply<G: StaticGraphViewOps>(&self, graph: G) -> Vec<String> {
194 #[cfg(feature = "search")]
195 {
196 let mut results = graph
197 .search_edges(self.0.clone(), 20, 0)
198 .unwrap()
199 .into_iter()
200 .map(|ev| format!("{}->{}", ev.src().name(), ev.dst().name()))
201 .collect::<Vec<_>>();
202 results.sort();
203 return results;
204 }
205 #[cfg(not(feature = "search"))]
206 Vec::<String>::new()
207 }
208}
209
210pub fn assert_filter_nodes_results(
211 init_graph: impl FnOnce(Graph) -> Graph,
212 transform: impl GraphTransformer,
213 filter: impl TryAsCompositeFilter + CreateFilter + Clone,
214 expected: &[&str],
215 variants: impl Into<Vec<TestGraphVariants>>,
216) {
217 assert_results(
218 init_graph,
219 |_graph: &Graph| (),
220 transform,
221 expected,
222 variants.into(),
223 FilterNodes(filter),
224 )
225}
226
227pub fn assert_select_nodes_results(
228 init_graph: impl FnOnce(Graph) -> Graph,
229 transform: impl GraphTransformer,
230 filter: impl TryAsCompositeFilter + CreateFilter + Clone,
231 expected: &[&str],
232 variants: impl Into<Vec<TestGraphVariants>>,
233) {
234 assert_results(
235 init_graph,
236 |_graph: &Graph| (),
237 transform,
238 expected,
239 variants.into(),
240 SelectNodes(filter),
241 )
242}
243
244fn assert_filter_err_contains<E>(err: GraphError, expected: E)
245where
246 E: AsRef<str>,
247{
248 match err {
249 GraphError::InvalidFilter(msg) => {
250 assert!(
251 msg.contains(expected.as_ref()),
252 "unexpected InvalidFilter message.\nexpected to contain: {}\nactual: {}",
253 expected.as_ref(),
254 msg
255 );
256 }
257 other => panic!("expected InvalidFilter, got: {other:?}"),
258 }
259}
260
261pub fn assert_filter_nodes_err(
262 init_graph: fn(Graph) -> Graph,
263 transform: impl GraphTransformer,
264 filter: impl TryAsCompositeFilter + CreateFilter + Clone,
265 expected: &str,
266 variants: impl Into<Vec<TestGraphVariants>>,
267) {
268 let graph = init_graph(Graph::new());
269 let variants = variants.into();
270
271 for v in variants {
272 match v {
273 TestGraphVariants::Graph => {
274 let graph = transform.apply(graph.clone());
275 let res = graph.filter(filter.clone());
276 assert!(res.is_err(), "expected error, filter was accepted");
277 assert_filter_err_contains(res.err().unwrap(), expected);
278 }
279 TestGraphVariants::PersistentGraph => {
280 let base = graph.persistent_graph();
281 let graph = transform.apply(base);
282 let res = graph.filter(filter.clone());
283 assert!(res.is_err(), "expected error, filter was accepted");
284 assert_filter_err_contains(res.err().unwrap(), expected);
285 }
286 TestGraphVariants::EventDiskGraph => {
287 #[cfg(feature = "storage")]
288 {
289 let tmp = TempDir::new().unwrap();
290 let graph = graph.persist_as_disk_graph(tmp.path()).unwrap();
291 let graph = transform.apply(graph);
292 let res = graph.filter(filter.clone());
293 assert!(res.is_err(), "expected error, filter was accepted");
294 assert_filter_err_contains(res.err().unwrap(), expected);
295 }
296 }
297 TestGraphVariants::PersistentDiskGraph => {
298 #[cfg(feature = "storage")]
299 {
300 let tmp = TempDir::new().unwrap();
301 let disk = DiskGraphStorage::from_graph(&graph, &tmp).unwrap();
302 let graph = disk.into_graph().persistent_graph();
303 let graph = transform.apply(graph);
304 let res = graph.filter(filter.clone());
305 assert!(res.is_err(), "expected error, filter was accepted");
306 assert_filter_err_contains(res.err().unwrap(), expected);
307 }
308 }
309 }
310 }
311}
312
313pub fn assert_filter_neighbours_results(
314 init_graph: impl FnOnce(Graph) -> Graph,
315 transform: impl GraphTransformer,
316 node_name: impl AsRef<str>,
317 direction: Direction,
318 filter: impl TryAsCompositeFilter + CreateFilter + Clone,
319 expected: &[&str],
320 variants: impl Into<Vec<TestGraphVariants>>,
321) {
322 assert_results(
323 init_graph,
324 |_graph: &Graph| (),
325 transform,
326 expected,
327 variants.into(),
328 FilterNeighbours(filter, node_name.as_ref().to_string(), direction),
329 )
330}
331
332pub fn assert_search_nodes_results(
333 init_graph: impl FnOnce(Graph) -> Graph,
334 transform: impl GraphTransformer,
335 filter: impl TryAsCompositeFilter + CreateFilter + Clone,
336 expected: &[&str],
337 variants: impl Into<Vec<TestGraphVariants>>,
338) {
339 #[cfg(feature = "search")]
340 {
341 assert_results(
342 init_graph,
343 |graph: &Graph| graph.create_index_in_ram().unwrap(),
344 transform,
345 expected,
346 variants.into(),
347 SearchNodes(filter),
348 )
349 }
350}
351
352pub fn assert_filter_edges_results(
353 init_graph: impl FnOnce(Graph) -> Graph,
354 transform: impl GraphTransformer,
355 filter: impl TryAsCompositeFilter + CreateFilter + Clone,
356 expected: &[&str],
357 variants: impl Into<Vec<TestGraphVariants>>,
358) {
359 assert_results(
360 init_graph,
361 |_graph: &Graph| (),
362 transform,
363 expected,
364 variants.into(),
365 FilterEdges(filter),
366 )
367}
368
369pub fn assert_select_edges_results(
370 init_graph: impl FnOnce(Graph) -> Graph,
371 transform: impl GraphTransformer,
372 filter: impl TryAsCompositeFilter + CreateFilter + Clone,
373 expected: &[&str],
374 variants: impl Into<Vec<TestGraphVariants>>,
375) {
376 assert_results(
377 init_graph,
378 |_graph: &Graph| (),
379 transform,
380 expected,
381 variants.into(),
382 SelectEdges(filter),
383 )
384}
385
386pub fn assert_search_edges_results(
387 init_graph: impl FnOnce(Graph) -> Graph,
388 transform: impl GraphTransformer,
389 filter: impl TryAsCompositeFilter + CreateFilter + Clone,
390 expected: &[&str],
391 variants: impl Into<Vec<TestGraphVariants>>,
392) {
393 #[cfg(feature = "search")]
394 {
395 assert_results(
396 init_graph,
397 |graph: &Graph| graph.create_index_in_ram().unwrap(),
398 transform,
399 expected,
400 variants.into(),
401 SearchEdges(filter),
402 )
403 }
404}
405
406fn assert_results(
407 init_graph: impl FnOnce(Graph) -> Graph,
408 pre_transform: impl Fn(&Graph) -> (),
409 transform: impl GraphTransformer,
410 expected: &[&str],
411 variants: Vec<TestGraphVariants>,
412 apply: impl ApplyFilter,
413) {
414 fn sorted<I, S>(iter: I) -> Vec<String>
415 where
416 I: IntoIterator<Item = S>,
417 S: AsRef<str>,
418 {
419 let mut v: Vec<String> = iter.into_iter().map(|s| s.as_ref().to_string()).collect();
420 v.sort();
421 v
422 }
423
424 let graph = init_graph(Graph::new());
425
426 let expected = sorted(expected.iter());
427
428 for v in variants {
429 match v {
430 TestGraphVariants::Graph => {
431 pre_transform(&graph);
432 let graph = transform.apply(graph.clone());
433 let result = sorted(apply.apply(graph));
434 assert_eq!(expected, result);
435 }
436 TestGraphVariants::PersistentGraph => {
437 pre_transform(&graph);
438 let base = graph.persistent_graph();
439 let graph = transform.apply(base);
440 let result = sorted(apply.apply(graph));
441 assert_eq!(expected, result);
442 }
443 TestGraphVariants::EventDiskGraph => {
444 #[cfg(feature = "storage")]
445 {
446 let tmp = TempDir::new().unwrap();
447 let graph = graph.persist_as_disk_graph(tmp.path()).unwrap();
448 pre_transform(&graph);
449 let graph = transform.apply(graph);
450 let result = sorted(apply.apply(graph));
451 assert_eq!(expected, result);
452 }
453 }
454 TestGraphVariants::PersistentDiskGraph => {
455 #[cfg(feature = "storage")]
456 {
457 let tmp = TempDir::new().unwrap();
458 let graph = DiskGraphStorage::from_graph(&graph, &tmp).unwrap();
459 let graph = graph.into_graph();
460 pre_transform(&graph);
461 let graph = graph.persistent_graph();
462 let graph = transform.apply(graph);
463 let result = sorted(apply.apply(graph));
464 assert_eq!(expected, result);
465 }
466 }
467 }
468 }
469}
470
471pub fn filter_nodes(graph: &Graph, filter: impl CreateFilter) -> Vec<String> {
472 let mut results = graph
473 .filter(filter)
474 .unwrap()
475 .nodes()
476 .iter()
477 .map(|n| n.name())
478 .collect::<Vec<_>>();
479 results.sort();
480 results
481}
482
483#[cfg(feature = "search")]
484pub fn search_nodes(graph: &Graph, filter: impl TryAsCompositeFilter) -> Vec<String> {
485 let mut results = graph
486 .search_nodes(filter, 10, 0)
487 .expect("Failed to search nodes")
488 .into_iter()
489 .map(|v| v.name())
490 .collect::<Vec<_>>();
491 results.sort();
492 results
493}
494
495pub fn filter_edges(graph: &Graph, filter: impl CreateFilter) -> Vec<String> {
496 let mut results = graph
497 .filter(filter)
498 .unwrap()
499 .edges()
500 .iter()
501 .map(|e| format!("{}->{}", e.src().name(), e.dst().name()))
502 .collect::<Vec<_>>();
503 results.sort();
504 results
505}
506
507#[cfg(feature = "search")]
508pub fn search_edges(graph: &Graph, filter: impl TryAsCompositeFilter) -> Vec<String> {
509 let mut results = graph
510 .search_edges(filter, 10, 0)
511 .expect("Failed to filter edges")
512 .into_iter()
513 .map(|e| format!("{}->{}", e.src().name(), e.dst().name()))
514 .collect::<Vec<_>>();
515 results.sort();
516 results
517}
518
519pub type EdgeRow = (u64, u64, i64, String, i64);
520
521pub fn assert_ok_or_missing_edges<T>(
522 edges: &[EdgeRow],
523 res: Result<T, GraphError>,
524 on_ok: impl FnOnce(T),
525) {
526 match res {
527 Ok(v) => on_ok(v),
528 Err(GraphError::PropertyMissingError(name)) => {
529 assert!(
530 edges.is_empty(),
531 "PropertyMissingError({name}) on non-empty graph"
532 );
533 }
534 Err(err) => panic!("Unexpected error from filter: {err:?}"),
535 }
536}
537
538pub fn assert_ok_or_missing_nodes<T>(
539 nodes: &[(u64, Option<String>, Option<i64>)],
540 res: Result<T, GraphError>,
541 on_ok: impl FnOnce(T),
542) {
543 match res {
544 Ok(v) => on_ok(v),
545
546 Err(GraphError::PropertyMissingError(name)) => {
547 let property_really_missing = match name.as_str() {
548 "int_prop" => nodes.iter().all(|(_, _, iv)| iv.is_none()),
549 "str_prop" => nodes.iter().all(|(_, sv, _)| sv.is_none()),
550 _ => panic!("Unexpected property {name}"),
551 };
552
553 assert!(
554 property_really_missing,
555 "PropertyMissingError({name}) but at least one node had that property"
556 );
557 }
558
559 Err(err) => panic!("Unexpected error from filter: {err:?}"),
560 }
561}