Skip to content

Commit 054e2ce

Browse files
rscarvalhorscarvalho
committed
Adding support for SdkTestingExecutor to run dynamic workflows
Addding a javadoc to SdkTestingExecutor Signed-off-by: Rodolfo Carvalho <rodolfo.ueg@gmail.com>
1 parent d8c898c commit 054e2ce

File tree

4 files changed

+237
-0
lines changed

4 files changed

+237
-0
lines changed

‎flytekit-testing/src/main/java/org/flyte/flytekit/testing/SdkDynamicWorkflowTaskDelegatingWorkflow.java

+44
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/*
2+
* Copyright 2025 Flyte Authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing,
11+
* software distributed under the License is distributed on an
12+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
13+
* KIND, either express or implied. See the License for the
14+
* specific language governing permissions and limitations
15+
* under the License.
16+
*/
17+
package org.flyte.flytekit.testing;
18+
19+
import java.util.Objects;
20+
import org.flyte.flytekit.SdkDynamicWorkflowTask;
21+
import org.flyte.flytekit.SdkType;
22+
import org.flyte.flytekit.SdkTypes;
23+
import org.flyte.flytekit.SdkWorkflow;
24+
import org.flyte.flytekit.SdkWorkflowBuilder;
25+
26+
public class SdkDynamicWorkflowTaskDelegatingWorkflow<InputT, OutputT>
27+
extends SdkWorkflow<Void, OutputT> {
28+
private final SdkDynamicWorkflowTask<InputT, OutputT> delegate;
29+
private final InputT input;
30+
31+
public SdkDynamicWorkflowTaskDelegatingWorkflow(
32+
SdkDynamicWorkflowTask<InputT, OutputT> delegate,
33+
InputT input,
34+
SdkType<OutputT> outputSdkType) {
35+
super(SdkTypes.nulls(), outputSdkType);
36+
this.delegate = Objects.requireNonNull(delegate, "delegate cannot be null");
37+
this.input = input;
38+
}
39+
40+
@Override
41+
protected OutputT expand(SdkWorkflowBuilder builder, Void ignored) {
42+
return delegate.run(builder, this.input);
43+
}
44+
}

‎flytekit-testing/src/main/java/org/flyte/flytekit/testing/SdkTestingExecutor.java

+36
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,42 @@ public static SdkTestingExecutor of(
116116
.build();
117117
}
118118

119+
/**
120+
* Creates a new {@link SdkTestingExecutor} for testing a {@link SdkDynamicWorkflowTask}.
121+
* This method wraps the given dynamic workflow task and its input into a delegating workflow,
122+
* allowing the task to be executed and tested in isolation.
123+
*
124+
* @param task the dynamic workflow task to test
125+
* @param input the input to the dynamic workflow task
126+
* @param outputType the expected output type of the dynamic workflow task
127+
* @param <InputT> the type of the input
128+
* @param <OutputT> the type of the output
129+
* @return a new {@link SdkTestingExecutor} instance
130+
*
131+
* <p>Example usage:
132+
* <pre>{@code
133+
* int expected = 6;
134+
*
135+
* SumIfEvenDynamicWorkflowTask.Output output =
136+
* SdkTestingExecutor.of(
137+
* new SumIfEvenDynamicWorkflowTask(),
138+
* SumIfEvenDynamicWorkflowTask.Input.create(of(2), of(4)),
139+
* JacksonSdkType.of(SumIfEvenDynamicWorkflowTask.Output.class))
140+
* .withTaskOutput(
141+
* new SumTask(),
142+
* SumTask.SumInput.create(of(2), of(4)),
143+
* SumTask.SumOutput.create(of(expected)))
144+
* .execute()
145+
* .getOutputAs(JacksonSdkType.of(SumIfEvenDynamicWorkflowTask.Output.class));
146+
*
147+
* assertEquals(expected, output.c().get());
148+
* }</pre>
149+
*/
150+
public static <InputT, OutputT> SdkTestingExecutor of(
151+
SdkDynamicWorkflowTask<InputT, OutputT> task, InputT input, SdkType<OutputT> outputType) {
152+
return of(new SdkDynamicWorkflowTaskDelegatingWorkflow<>(task, input, outputType));
153+
}
154+
119155
@AutoValue
120156
public abstract static class Result {
121157
abstract Map<String, Literal> literalMap();

‎flytekit-testing/src/test/java/org/flyte/flytekit/testing/SdkDynamicWorkflowTaskDelegatingWorkflowTest.java

+79
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
/*
2+
* Copyright 2025 Flyte Authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing,
11+
* software distributed under the License is distributed on an
12+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
13+
* KIND, either express or implied. See the License for the
14+
* specific language governing permissions and limitations
15+
* under the License.
16+
*/
17+
package org.flyte.flytekit.testing;
18+
19+
import static org.flyte.flytekit.SdkBindingDataFactory.of;
20+
import static org.junit.jupiter.api.Assertions.*;
21+
22+
import org.flyte.flytekit.jackson.JacksonSdkType;
23+
import org.junit.jupiter.api.Test;
24+
25+
public class SdkDynamicWorkflowTaskDelegatingWorkflowTest {
26+
@Test
27+
public void testDelegatingWorkflow_EvenValues() {
28+
int expected = 6;
29+
30+
SumIfEvenDynamicWorkflowTask.Output output =
31+
SdkTestingExecutor.of(
32+
new SumIfEvenDynamicWorkflowTask(),
33+
SumIfEvenDynamicWorkflowTask.Input.create(of(2), of(4)),
34+
JacksonSdkType.of(SumIfEvenDynamicWorkflowTask.Output.class))
35+
.withTaskOutput(
36+
new SumTask(),
37+
SumTask.SumInput.create(of(2), of(4)),
38+
SumTask.SumOutput.create(of(expected)))
39+
.execute()
40+
.getOutputAs(JacksonSdkType.of(SumIfEvenDynamicWorkflowTask.Output.class));
41+
assertEquals(expected, output.c().get());
42+
}
43+
44+
@Test
45+
public void testDelegatingWorkflow_OddA() {
46+
int expected = 0;
47+
48+
SumIfEvenDynamicWorkflowTask.Output output =
49+
SdkTestingExecutor.of(
50+
new SumIfEvenDynamicWorkflowTask(),
51+
SumIfEvenDynamicWorkflowTask.Input.create(of(1), of(4)),
52+
JacksonSdkType.of(SumIfEvenDynamicWorkflowTask.Output.class))
53+
.withTaskOutput(
54+
new SumTask(),
55+
SumTask.SumInput.create(of(0), of(0)),
56+
SumTask.SumOutput.create(of(expected)))
57+
.execute()
58+
.getOutputAs(JacksonSdkType.of(SumIfEvenDynamicWorkflowTask.Output.class));
59+
assertEquals(expected, output.c().get());
60+
}
61+
62+
@Test
63+
public void testDelegatingWorkflow_OddB() {
64+
int expected = 0;
65+
66+
SumIfEvenDynamicWorkflowTask.Output output =
67+
SdkTestingExecutor.of(
68+
new SumIfEvenDynamicWorkflowTask(),
69+
SumIfEvenDynamicWorkflowTask.Input.create(of(2), of(5)),
70+
JacksonSdkType.of(SumIfEvenDynamicWorkflowTask.Output.class))
71+
.withTaskOutput(
72+
new SumTask(),
73+
SumTask.SumInput.create(of(0), of(0)),
74+
SumTask.SumOutput.create(of(expected)))
75+
.execute()
76+
.getOutputAs(JacksonSdkType.of(SumIfEvenDynamicWorkflowTask.Output.class));
77+
assertEquals(expected, output.c().get());
78+
}
79+
}

‎flytekit-testing/src/test/java/org/flyte/flytekit/testing/SumIfEvenDynamicWorkflowTask.java

+78
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* Copyright 2025 Flyte Authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing,
11+
* software distributed under the License is distributed on an
12+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
13+
* KIND, either express or implied. See the License for the
14+
* specific language governing permissions and limitations
15+
* under the License.
16+
*/
17+
package org.flyte.flytekit.testing;
18+
19+
import static org.flyte.flytekit.SdkBindingDataFactory.of;
20+
21+
import com.google.auto.service.AutoService;
22+
import com.google.auto.value.AutoValue;
23+
import org.flyte.flytekit.*;
24+
import org.flyte.flytekit.jackson.JacksonSdkType;
25+
26+
@AutoService(SumIfEvenDynamicWorkflowTask.class)
27+
public class SumIfEvenDynamicWorkflowTask
28+
extends SdkDynamicWorkflowTask<
29+
SumIfEvenDynamicWorkflowTask.Input, SumIfEvenDynamicWorkflowTask.Output> {
30+
@AutoValue
31+
public abstract static class Input {
32+
33+
abstract SdkBindingData<Long> a();
34+
35+
abstract SdkBindingData<Long> b();
36+
37+
static Input create(SdkBindingData<Long> a, SdkBindingData<Long> b) {
38+
return new AutoValue_SumIfEvenDynamicWorkflowTask_Input(a, b);
39+
}
40+
}
41+
42+
@AutoValue
43+
public abstract static class Output {
44+
45+
abstract SdkBindingData<Long> c();
46+
47+
static Output create(SdkBindingData<Long> c) {
48+
return new AutoValue_SumIfEvenDynamicWorkflowTask_Output(c);
49+
}
50+
}
51+
52+
public SumIfEvenDynamicWorkflowTask() {
53+
super(JacksonSdkType.of(Input.class), JacksonSdkType.of(Output.class));
54+
}
55+
56+
@Override
57+
public Output run(SdkWorkflowBuilder builder, Input input) {
58+
/*
59+
* This is to demonstrate that we can use concrete values in the dynamic workflow task
60+
*/
61+
long aConcreteValue = input.a().get();
62+
long bConcreteValue = input.b().get();
63+
64+
SumTask.SumOutput outputs =
65+
builder
66+
.apply(
67+
SdkConditions.when(
68+
"is-even",
69+
SdkConditions.isTrue(
70+
of(aConcreteValue % 2 == 0 && bConcreteValue % 2 == 0)),
71+
new SumTask(),
72+
SumTask.SumInput.create(input.a(), input.b()))
73+
.otherwise("is-odd", new SumTask(), SumTask.SumInput.create(of(0L), of(0L))))
74+
.getOutputs();
75+
76+
return Output.create(outputs.c());
77+
}
78+
}

0 commit comments

Comments
 (0)