graph_api_test/steps/
control_flow.rs

1use crate::{Edge, Vertex, assert_elements_eq, populate_graph};
2use graph_api_lib::{EdgeReference, EdgeSearch, Graph, VertexReference};
3use std::ops::ControlFlow;
4
5pub fn test_vertices_control_flow<T>(graph: &mut T)
6where
7    T: Graph<Vertex = Vertex, Edge = Edge>,
8{
9    let refs = populate_graph(graph);
10
11    // Test Continue(Some) / Continue(None) - Should work like filter
12    let collected = graph
13        .walk()
14        .vertices_by_id(vec![refs.bryn, refs.julia, refs.graph_api])
15        .control_flow(|vertex, _| match vertex.weight() {
16            Vertex::Person { .. } => ControlFlow::Continue(Some(vertex)),
17            _ => ControlFlow::Continue(None), // Skip non-Person vertices
18        })
19        .collect::<Vec<_>>();
20    assert_elements_eq!(graph, collected, vec![refs.bryn, refs.julia]);
21
22    // Test Break with Some - Should stop traversal after finding Julia but include it
23    let collected = graph
24        .walk()
25        .vertices_by_id(vec![refs.bryn, refs.julia, refs.graph_api])
26        .control_flow(|vertex, _| match vertex.weight() {
27            Vertex::Person { name, .. } if name == "Bryn" => ControlFlow::Continue(Some(vertex)),
28            Vertex::Person { name, .. } if name == "Julia" => ControlFlow::Break(Some(vertex)),
29            _ => ControlFlow::Continue(None),
30        })
31        .collect::<Vec<_>>();
32    assert_elements_eq!(graph, collected, vec![refs.bryn, refs.julia]);
33
34    // Test Break with None - Should stop traversal after finding Julia without including it
35    let collected = graph
36        .walk()
37        .vertices_by_id(vec![refs.bryn, refs.julia, refs.graph_api])
38        .control_flow(|vertex, _| match vertex.weight() {
39            Vertex::Person { name, .. } if name == "Bryn" => ControlFlow::Continue(Some(vertex)),
40            Vertex::Person { name, .. } if name == "Julia" => ControlFlow::Break(None),
41            _ => ControlFlow::Continue(Some(vertex)),
42        })
43        .collect::<Vec<_>>();
44    assert_elements_eq!(graph, collected, vec![refs.bryn]);
45
46    // Test context mutation with control flow
47    let collected = graph
48        .walk()
49        .push_context(0)
50        .vertices_by_id(vec![refs.bryn, refs.julia, refs.graph_api])
51        .control_flow(|vertex, ctx| {
52            *ctx += 1;
53            if *ctx >= 2 {
54                ControlFlow::Break(None)
55            } else {
56                ControlFlow::Continue(Some(vertex))
57            }
58        })
59        .collect::<Vec<_>>();
60    assert_elements_eq!(graph, collected, vec![refs.bryn]);
61}
62
63pub fn test_edges_control_flow<T>(graph: &mut T)
64where
65    T: Graph<Vertex = Vertex, Edge = Edge>,
66{
67    let refs = populate_graph(graph);
68
69    // As we don't know what order edges are returned we rely on context.
70    let collected = graph
71        .walk()
72        .push_context(0)
73        .vertices_by_id(vec![refs.bryn, refs.graph_api])
74        .edges(EdgeSearch::scan())
75        .control_flow(|edge, ctx| {
76            *ctx += 1;
77            if *ctx >= 2 {
78                ControlFlow::Break(None)
79            } else {
80                ControlFlow::Continue(Some(edge))
81            }
82        })
83        .count();
84    // Only the first edge should be processed
85    assert_eq!(collected, 1);
86
87    // Test skipping certain edges with None
88    let collected = graph
89        .walk()
90        .vertices_by_id(vec![refs.bryn])
91        .edges(EdgeSearch::scan())
92        .control_flow(|edge, _| {
93            if let Edge::Knows { .. } = edge.weight() {
94                ControlFlow::Continue(Some(edge))
95            } else {
96                ControlFlow::Continue(None) // Skip non-Knows edges
97            }
98        })
99        .collect::<Vec<_>>();
100
101    // We should only collect "knows" edges
102    for edge_id in &collected {
103        if let Some(edge) = graph.edge(*edge_id) {
104            assert!(matches!(edge.weight(), Edge::Knows { .. }));
105        }
106    }
107}
graph_api_test/steps/control_flow.rs

graph_api_test/steps/
control_flow.rs
