From 4721491878231bd9b7a448b2a7c8d25be6798654 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Fri, 16 Oct 2026 02:59:43 -0700 Subject: [PATCH] feat: add a builder for workflow graph edges PiperOrigin-RevId: 1000728765 --- .../com/google/adk/workflow/EdgesBuilder.java | 368 ++++++++++++++++++ .../com/google/adk/workflow/Workflow.java | 13 + .../google/adk/workflow/EdgesBuilderTest.java | 361 +++++++++++++++++ .../com/google/adk/workflow/WorkflowTest.java | 21 + 4 files changed, 763 insertions(+) create mode 100644 core/src/main/java/com/google/adk/workflow/EdgesBuilder.java create mode 100644 core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java diff --git a/core/src/main/java/com/google/adk/workflow/EdgesBuilder.java b/core/src/main/java/com/google/adk/workflow/EdgesBuilder.java new file mode 100644 index 000000000..ea5081142 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/EdgesBuilder.java @@ -0,0 +1,368 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.workflow; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; + +import com.google.adk.annotations.Experimental; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.function.Consumer; + +/** + * A builder for the edges of a workflow graph, declared as chains of calls that start at {@link + * #from} or {@link #fromStart}. Create one with {@code new EdgesBuilder()} and call {@link #build}, + * or let {@link Workflow.Builder#edges(Consumer)} do both. A node is identified by its instance, so + * create each node once. + * + * 
{@code
+ * Workflow support =
+ *     Workflow.builder()
+ *         .name("support_triage")
+ *         .edges(
+ *             g -> {
+ *               g.fromStart()
+ *                   .chain(classify, category)
+ *                   .route(
+ *                       r -> {
+ *                         r.on("billing").then(billing);
+ *                         r.otherwise(summarize);
+ *                       });
+ *               g.from(billing).then(summarize);
+ *             })
+ *         .build();
+ * }
+ */ +@Experimental +public final class EdgesBuilder { + + private final List edges = new ArrayList<>(); + + /** Starts a chain of edges at {@code node}: {@code g.from(a).then(b).then(c)}. */ + public EdgeChain from(Node node) { + return new EdgeChain(this, checkNotNull(node)); + } + + /** + * Starts a chain at {@code nodes} (such as a computed list) so they can join without a fan-out, + * like {@link #from(Node, Node, Node...)}. + */ + public NodeGroupChain from(Collection nodes) { + return new NodeGroupChain(this, ImmutableList.copyOf(nodes)); + } + + /** + * Starts a chain at {@code first}, {@code second} and each node in {@code rest} so they can join + * without a fan-out: {@code g.from(a, b).joinTo(join)}. + */ + public NodeGroupChain from(Node first, Node second, Node... rest) { + return from(Lists.asList(first, second, rest)); + } + + /** Starts a chain of edges at {@link Start}: {@code g.fromStart().then(a)}. */ + public EdgeChain fromStart() { + return from(Start.INSTANCE); + } + + /** Returns the edges added so far, in the order they were added. */ + public ImmutableList build() { + return ImmutableList.copyOf(edges); + } + + void add(Edge edge) { + edges.add(edge); + } + + void addAll(List newEdges) { + edges.addAll(newEdges); + } + + /** Returns an edge from {@code from} to each of {@code targets}, failing if a target is null. */ + static ImmutableList edgesTo( + Node from, Collection targets, ImmutableList routes) { + return targets.stream().map(to -> new Edge(from, to, routes)).collect(toImmutableList()); + } + + /** + * A chain of edges positioned at a node, started by {@link EdgesBuilder#from(Node)}, {@link + * EdgesBuilder#fromStart} or {@link NodeGroupChain#joinTo}. {@link #then(Node)} and {@link + * #chain} add edges and continue from the last node they connect, a fan-out continues from its + * targets as a group, and a routing map ends the chain. + */ + @Experimental + public static final class EdgeChain { + + private final EdgesBuilder builder; + private final Node node; + + EdgeChain(EdgesBuilder builder, Node node) { + this.builder = builder; + this.node = node; + } + + /** Connects the current node to {@code next} and continues from {@code next}. */ + @CanIgnoreReturnValue + public EdgeChain then(Node next) { + builder.add(new Edge(node, next)); + return new EdgeChain(builder, next); + } + + /** + * Fans out by connecting the current node to each of {@code targets} (such as a computed list) + * so they run concurrently, and continues from them as a group that {@link + * NodeGroupChain#joinTo} can join. An empty collection adds no edges. + */ + @CanIgnoreReturnValue + public NodeGroupChain then(Collection targets) { + ImmutableList nodes = ImmutableList.copyOf(targets); + builder.addAll(edgesTo(node, nodes, ImmutableList.of())); + return new NodeGroupChain(builder, nodes); + } + + /** + * Fans out by connecting the current node to {@code first}, {@code second} and each node in + * {@code rest}, like {@link #then(Collection)}: {@code g.from(a).then(b, c).joinTo(join)}. + */ + @CanIgnoreReturnValue + public NodeGroupChain then(Node first, Node second, Node... rest) { + return then(Lists.asList(first, second, rest)); + } + + /** + * Connects the current node to {@code first}, {@code second} and each node in {@code rest} in + * order, and continues from the last one: {@code g.fromStart().chain(a, b)} is short for {@code + * g.fromStart().then(a).then(b)}. + */ + @CanIgnoreReturnValue + public EdgeChain chain(Node first, Node second, Node... rest) { + // Copying rejects a null node before any edge is added. + ImmutableList nodes = ImmutableList.copyOf(Lists.asList(first, second, rest)); + EdgeChain chain = this; + for (Node next : nodes) { + chain = chain.then(next); + } + return chain; + } + + /** + * Adds the current node's routing map, which {@code routeMapDefinition} fills, and ends the + * chain. + * + * @throws IllegalArgumentException if the routing map has no entries + * @throws IllegalStateException if an entry started by {@link RouteMap#on} has no target + */ + public void route(Consumer routeMapDefinition) { + RouteMap routeMap = new RouteMap(node); + routeMapDefinition.accept(routeMap); + ImmutableList routeEdges = routeMap.build(); + checkArgument( + !routeEdges.isEmpty(), + "The routing map of node '%s' needs at least one entry.", + node.name()); + builder.addAll(routeEdges); + } + } + + /** + * A chain of edges positioned at a node group from a fan-out, {@link EdgesBuilder#from(Node, + * Node, Node...)}, or {@link EdgesBuilder#from(Collection)}, which {@link #joinTo} connects to a + * {@link JoinNode}. + */ + @Experimental + public static final class NodeGroupChain { + + private final EdgesBuilder builder; + private final ImmutableList nodes; + + NodeGroupChain(EdgesBuilder builder, ImmutableList nodes) { + this.builder = builder; + this.nodes = nodes; + } + + /** + * Connects each node in the group to {@code join}, which runs once all of them complete, and + * continues from {@code join}. + */ + @CanIgnoreReturnValue + public EdgeChain joinTo(JoinNode join) { + checkNotNull(join); + for (Node node : nodes) { + builder.add(new Edge(node, join)); + } + return new EdgeChain(builder, join); + } + } + + /** + * A routing map configured in {@link EdgeChain#route} that maps each route the source node may + * emit to the nodes it leads to, written {@code r.on("billing").then(billing)}. Mapping a route + * to several nodes fans out on that route. + */ + @Experimental + public static final class RouteMap { + + private final Node from; + private final ImmutableList.Builder edges = ImmutableList.builder(); + private final List pendingEntries = new ArrayList<>(); + private boolean built = false; + + RouteMap(Node from) { + this.from = from; + } + + /** + * Starts an entry for {@code route} as a {@link Route.Tag}; {@link RouteMapEntry#then} + * completes it. + */ + public RouteMapEntry on(String route) { + return on(new Route.Tag(route)); + } + + /** + * Starts an entry for {@code route} as a {@link Route.Flag}; {@link RouteMapEntry#then} + * completes it. + */ + public RouteMapEntry on(boolean route) { + return on(new Route.Flag(route)); + } + + /** + * Starts an entry for {@code route}, such as a {@link Route.Num}; {@link RouteMapEntry#then} + * completes it. + */ + public RouteMapEntry on(Route route) { + return start(ImmutableList.of(route)); + } + + /** + * Starts an entry whose edges fire once when the source node emits any of {@code routes}, built + * with {@link #anyOf}; {@link RouteMapEntry#then} completes it. + */ + public RouteMapEntry on(RouteSet routes) { + return start(routes.routes()); + } + + /** + * Leads to each target when no other route of the source node matches, including when it emits + * none. + */ + public void otherwise(Node first, Node... rest) { + add(ImmutableList.of(Route.Default.INSTANCE), Lists.asList(first, rest)); + } + + /** Groups string routes as {@link Route.Tag}s for {@link #on(RouteSet)}. */ + public RouteSet anyOf(String first, String second, String... rest) { + return new RouteSet(Lists.transform(Lists.asList(first, second, rest), Route.Tag::new)); + } + + /** Groups routes, such as {@link Route.Num}s, for {@link #on(RouteSet)}. */ + public RouteSet anyOf(Route first, Route second, Route... rest) { + return new RouteSet(Lists.asList(first, second, rest)); + } + + /** + * Returns the map's edges and closes the map, so an entry started or completed after {@link + * EdgeChain#route} returns fails instead of being lost. + * + * @throws IllegalStateException if an entry started by {@code on} has no target + */ + ImmutableList build() { + built = true; + checkState( + pendingEntries.isEmpty(), + "on(...) in the routing map of node '%s' has no target; complete it with then(node).", + from.name()); + return edges.build(); + } + + /** Adds an edge for {@code routes} to each of {@code targets} and completes {@code entry}. */ + void complete(RouteMapEntry entry, ImmutableList routes, List targets) { + add(routes, targets); + pendingEntries.remove(entry); + } + + private RouteMapEntry start(ImmutableList routes) { + checkNotBuilt(); + RouteMapEntry entry = new RouteMapEntry(this, routes); + pendingEntries.add(entry); + return entry; + } + + private void add(ImmutableList routes, List targets) { + checkNotBuilt(); + edges.addAll(edgesTo(from, targets, routes)); + } + + private void checkNotBuilt() { + checkState( + !built, + "The routing map of node '%s' was already built; add its entries inside route(...).", + from.name()); + } + } + + /** + * A routing-map entry started by {@link RouteMap#on} whose routes have no target yet; {@link + * #then} completes it. + */ + @Experimental + public static final class RouteMapEntry { + + private final RouteMap routeMap; + private final ImmutableList routes; + + RouteMapEntry(RouteMap routeMap, ImmutableList routes) { + this.routeMap = routeMap; + this.routes = routes; + } + + /** + * Leads to {@code first} and each node in {@code rest}, with one edge to each, when the source + * node emits any of this entry's routes. + */ + public void then(Node first, Node... rest) { + routeMap.complete(this, routes, Lists.asList(first, rest)); + } + } + + /** + * A set of routes that select a single edge, which fires once when its source emits any of them; + * separate edges to the same target would fire once each. {@link RouteMap#anyOf} builds one and + * drops duplicate routes. + */ + @Experimental + public static final class RouteSet { + + private final ImmutableList routes; + + RouteSet(List routes) { + this.routes = ImmutableSet.copyOf(routes).asList(); + } + + ImmutableList routes() { + return routes; + } + } +} diff --git a/core/src/main/java/com/google/adk/workflow/Workflow.java b/core/src/main/java/com/google/adk/workflow/Workflow.java index a486c0b4f..4d2dc00db 100644 --- a/core/src/main/java/com/google/adk/workflow/Workflow.java +++ b/core/src/main/java/com/google/adk/workflow/Workflow.java @@ -28,6 +28,7 @@ import io.reactivex.rxjava3.core.Flowable; import java.util.List; import java.util.Optional; +import java.util.function.Consumer; import org.jspecify.annotations.Nullable; /** @@ -118,6 +119,18 @@ public Builder edges(Edge... edges) { return edges(ImmutableList.copyOf(edges)); } + /** + * Sets the edges that {@code graphDefinition} adds to a new {@link EdgesBuilder}. + * + *

For example: {@code .edges(g -> g.fromStart().chain(a, b))}. + */ + @CanIgnoreReturnValue + public Builder edges(Consumer graphDefinition) { + EdgesBuilder builder = new EdgesBuilder(); + graphDefinition.accept(builder); + return edges(builder.build()); + } + @CanIgnoreReturnValue public Builder description(String description) { this.description = checkNotNull(description); diff --git a/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java b/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java new file mode 100644 index 000000000..248dcf4f7 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java @@ -0,0 +1,361 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.workflow; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.common.collect.ImmutableList; +import java.util.Arrays; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class EdgesBuilderTest { + + private final EdgesBuilder builder = new EdgesBuilder(); + private final Node a = new StubNode("a"); + private final Node b = new StubNode("b"); + private final Node c = new StubNode("c"); + private final Node d = new StubNode("d"); + private final JoinNode join = JoinNode.builder().name("join").build(); + + @Test + public void build_withoutEdges_returnsAnEmptyList() { + assertThat(builder.build()).isEmpty(); + } + + @Test + public void fromStart_chainOfThens_connectsEachNodeToTheNext() { + builder.fromStart().then(a).then(b); + + assertThat(builder.build()) + .containsExactly(new Edge(Start.INSTANCE, a), new Edge(a, b)) + .inOrder(); + } + + @Test + public void from_startsTheChainAtTheNode() { + builder.from(a).then(b); + + assertThat(builder.build()).containsExactly(new Edge(a, b)); + } + + @Test + public void from_rejectsNull() { + assertThrows(NullPointerException.class, () -> builder.from((Node) null)); + } + + @Test + public void chain_severalNodes_connectsEachNodeToTheNext() { + builder.fromStart().chain(a, b, c); + + assertThat(builder.build()) + .containsExactly(new Edge(Start.INSTANCE, a), new Edge(a, b), new Edge(b, c)) + .inOrder(); + } + + @Test + public void chain_twoNodes_continuesFromTheLastOne() { + builder.from(a).chain(b, c).then(d); + + assertThat(builder.build()) + .containsExactly(new Edge(a, b), new Edge(b, c), new Edge(c, d)) + .inOrder(); + } + + @Test + public void chain_withANullNode_throwsAndAddsNoEdges() { + EdgesBuilder.EdgeChain chain = builder.from(a); + + assertThrows(NullPointerException.class, () -> chain.chain(b, c, (Node) null)); + assertThat(builder.build()).isEmpty(); + } + + @Test + public void then_collection_fansOutToEach() { + builder.fromStart().then(a).then(ImmutableList.of(b, c, d)); + + assertThat(builder.build()) + .containsExactly( + new Edge(Start.INSTANCE, a), new Edge(a, b), new Edge(a, c), new Edge(a, d)) + .inOrder(); + } + + @Test + public void then_emptyCollection_addsNoEdges() { + builder.fromStart().then(a).then(ImmutableList.of()); + + assertThat(builder.build()).containsExactly(new Edge(Start.INSTANCE, a)); + } + + @Test + public void then_collectionWithANull_throwsAndAddsNoEdges() { + EdgesBuilder.EdgeChain chain = builder.from(a); + + assertThrows(NullPointerException.class, () -> chain.then(Arrays.asList(b, null))); + assertThat(builder.build()).isEmpty(); + } + + @Test + public void then_collection_continuesAsAGroupThatJoins() { + builder.fromStart().then(ImmutableList.of(b, c)).joinTo(join).then(d); + + assertThat(builder.build()) + .containsExactly( + new Edge(Start.INSTANCE, b), + new Edge(Start.INSTANCE, c), + new Edge(b, join), + new Edge(c, join), + new Edge(join, d)) + .inOrder(); + } + + @Test + public void then_severalNodes_fansOutToEachAndContinuesAsAGroup() { + builder.from(a).then(b, c, d).joinTo(join); + + assertThat(builder.build()) + .containsExactly( + new Edge(a, b), + new Edge(a, c), + new Edge(a, d), + new Edge(b, join), + new Edge(c, join), + new Edge(d, join)) + .inOrder(); + } + + @Test + public void then_severalNodesWithANull_throwsAndAddsNoEdges() { + EdgesBuilder.EdgeChain chain = builder.from(a); + + assertThrows(NullPointerException.class, () -> chain.then(b, c, (Node) null)); + assertThat(builder.build()).isEmpty(); + } + + @Test + public void from_collection_joinsWithoutAFanOut() { + builder.from(ImmutableList.of(b, c)).joinTo(join).then(d); + + assertThat(builder.build()) + .containsExactly(new Edge(b, join), new Edge(c, join), new Edge(join, d)) + .inOrder(); + } + + @Test + public void from_collectionWithANull_throws() { + assertThrows(NullPointerException.class, () -> builder.from(Arrays.asList(b, null))); + } + + @Test + public void from_severalNodes_joinsWithoutAFanOut() { + builder.from(a, b, c).joinTo(join).then(d); + + assertThat(builder.build()) + .containsExactly(new Edge(a, join), new Edge(b, join), new Edge(c, join), new Edge(join, d)) + .inOrder(); + } + + @Test + public void from_severalNodesWithANull_throws() { + assertThrows(NullPointerException.class, () -> builder.from(a, b, (Node) null)); + } + + @Test + public void joinTo_emptyGroup_addsNoEdgesAndContinuesFromTheJoin() { + builder.from(ImmutableList.of()).joinTo(join).then(d); + + assertThat(builder.build()).containsExactly(new Edge(join, d)); + } + + @Test + public void joinTo_rejectsNull() { + // An empty group adds no edge that would reject the null join on its own. + EdgesBuilder.NodeGroupChain group = builder.from(ImmutableList.of()); + + assertThrows(NullPointerException.class, () -> group.joinTo(null)); + } + + @Test + public void route_everyRouteKind_routesFromTheCurrentNodeToEachTarget() { + ImmutableList anyRoute = ImmutableList.of(new Route.Tag("x"), new Route.Tag("y")); + + builder + .from(a) + .route( + r -> { + r.on("tag").then(b, c); + r.on(true).then(b, c); + r.on(new Route.Num(7)).then(b, c); + r.on(r.anyOf("x", "y")).then(b, c); + }); + + assertThat(builder.build()) + .containsExactly( + new Edge(a, b, new Route.Tag("tag")), + new Edge(a, c, new Route.Tag("tag")), + new Edge(a, b, new Route.Flag(true)), + new Edge(a, c, new Route.Flag(true)), + new Edge(a, b, new Route.Num(7)), + new Edge(a, c, new Route.Num(7)), + new Edge(a, b, anyRoute), + new Edge(a, c, anyRoute)) + .inOrder(); + } + + @Test + public void route_otherwiseWithSeveralTargets_addsADefaultEdgeToEach() { + builder + .from(a) + .route( + r -> { + r.on("x").then(b); + r.otherwise(b, c, d); + }); + + assertThat(builder.build()) + .containsExactly( + new Edge(a, b, new Route.Tag("x")), + new Edge(a, b, Route.Default.INSTANCE), + new Edge(a, c, Route.Default.INSTANCE), + new Edge(a, d, Route.Default.INSTANCE)) + .inOrder(); + } + + @Test + public void route_emptyRoutingMap_throws() { + EdgesBuilder.EdgeChain chain = builder.from(a); + + IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, () -> chain.route(r -> {})); + + assertThat(e).hasMessageThat().contains("routing map of node 'a'"); + } + + @Test + public void route_entryWithoutATarget_throwsAndAddsNoEdges() { + EdgesBuilder.EdgeChain chain = builder.from(a); + + IllegalStateException e = + assertThrows( + IllegalStateException.class, + () -> + chain.route( + r -> { + r.on("x").then(b); + EdgesBuilder.RouteMapEntry unused = r.on("y"); + })); + + assertThat(e).hasMessageThat().contains("routing map of node 'a' has no target"); + assertThat(builder.build()).isEmpty(); + } + + @Test + public void route_routeMapDefinitionThatThrows_addsNoEdges() { + EdgesBuilder.EdgeChain chain = builder.from(a); + + assertThrows( + NullPointerException.class, + () -> + chain.route( + r -> { + r.on("ok").then(b); + r.on("bad").then(c, (Node) null); + })); + assertThat(builder.build()).isEmpty(); + } + + @Test + public void route_entryStartedAfterTheRouteReturns_throws() { + AtomicReference escaped = new AtomicReference<>(); + builder + .from(a) + .route( + r -> { + escaped.set(r); + r.on("x").then(b); + }); + + IllegalStateException e = + assertThrows(IllegalStateException.class, () -> escaped.get().on("late")); + + assertThat(e).hasMessageThat().contains("routing map of node 'a' was already built"); + assertThat(builder.build()).containsExactly(new Edge(a, b, new Route.Tag("x"))); + } + + @Test + public void route_entryCompletedAfterTheRouteReturns_throws() { + AtomicReference escaped = new AtomicReference<>(); + builder + .from(a) + .route( + r -> { + escaped.set(r.on("x")); + escaped.get().then(b); + }); + + IllegalStateException e = + assertThrows(IllegalStateException.class, () -> escaped.get().then(c)); + + assertThat(e).hasMessageThat().contains("routing map of node 'a' was already built"); + assertThat(builder.build()).containsExactly(new Edge(a, b, new Route.Tag("x"))); + } + + @Test + public void anyOf_everyOverload_dropsRepeatedRoutes() { + EdgesBuilder.RouteMap r = new EdgesBuilder.RouteMap(a); + Route.Tag p = new Route.Tag("p"); + Route.Num seven = new Route.Num(7); + + assertThat(r.anyOf("x", "y", "z", "x").routes()) + .containsExactly(new Route.Tag("x"), new Route.Tag("y"), new Route.Tag("z")) + .inOrder(); + assertThat(r.anyOf(p, new Route.Flag(true), seven, p).routes()) + .containsExactly(p, new Route.Flag(true), seven) + .inOrder(); + } + + @Test + public void build_documentedExample_buildsItsGraph() { + Node classify = new StubNode("classifier"); + Node category = new StubNode("category"); + Node billing = new StubNode("billing"); + Node summarize = new StubNode("summary"); + + builder + .fromStart() + .chain(classify, category) + .route( + r -> { + r.on("billing").then(billing); + r.otherwise(summarize); + }); + builder.from(billing).then(summarize); + + assertThat(builder.build()) + .containsExactly( + new Edge(Start.INSTANCE, classify), + new Edge(classify, category), + new Edge(category, billing, new Route.Tag("billing")), + new Edge(category, summarize, Route.Default.INSTANCE), + new Edge(billing, summarize)) + .inOrder(); + } +} diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowTest.java b/core/src/test/java/com/google/adk/workflow/WorkflowTest.java index 457e6b47d..cf1f90aa2 100644 --- a/core/src/test/java/com/google/adk/workflow/WorkflowTest.java +++ b/core/src/test/java/com/google/adk/workflow/WorkflowTest.java @@ -25,6 +25,7 @@ import com.google.genai.types.Schema; import java.time.Duration; import java.util.List; +import java.util.function.Consumer; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -111,6 +112,19 @@ public void edges_withVarargs_setsTheEdgesInOrder() { assertThat(workflow.edges()).containsExactly(first, second).inOrder(); } + @Test + public void edges_withAConsumer_setsTheEdgesItAdds() { + Node a = new StubNode("a"); + Node b = new StubNode("b"); + + Workflow workflow = + Workflow.builder().name("flow").edges(g -> g.fromStart().then(a).then(b)).build(); + + assertThat(workflow.edges()) + .containsExactly(new Edge(Start.INSTANCE, a), new Edge(a, b)) + .inOrder(); + } + @Test public void name_rejectsNull() { Workflow.Builder builder = Workflow.builder(); @@ -139,6 +153,13 @@ public void edges_rejectsANullList() { assertThrows(NullPointerException.class, () -> builder.edges((List) null)); } + @Test + public void edges_rejectsANullConsumer() { + Workflow.Builder builder = Workflow.builder(); + + assertThrows(NullPointerException.class, () -> builder.edges((Consumer) null)); + } + @Test public void maxConcurrency_acceptsOne() { Workflow workflow = Workflow.builder().name("flow").maxConcurrency(1).build();