2828import io .temporal .worker .Worker ;
2929import io .temporal .worker .WorkerFactoryOptions ;
3030import io .temporal .worker .WorkerOptions ;
31+ import io .temporal .workflow .DynamicWorkflow ;
3132import java .lang .reflect .Constructor ;
33+ import java .lang .reflect .Parameter ;
3234import java .time .Instant ;
3335import java .util .HashSet ;
3436import java .util .Set ;
4042import org .junit .jupiter .api .extension .ParameterResolutionException ;
4143import org .junit .jupiter .api .extension .ParameterResolver ;
4244import org .junit .jupiter .api .extension .TestWatcher ;
45+ import org .junit .platform .commons .support .AnnotationSupport ;
4346
4447/**
4548 * JUnit Jupiter extension that simplifies testing of Temporal workflows.
@@ -77,7 +80,7 @@ public class TestWorkflowExtension
7780
7881 private static final String TEST_ENVIRONMENT_KEY = "testEnvironment" ;
7982 private static final String WORKER_KEY = "worker" ;
80- private static final String TASK_QUEUE_KEY = "taskQueue " ;
83+ private static final String WORKFLOW_OPTIONS_KEY = "workflowOptions " ;
8184
8285 private final WorkerOptions workerOptions ;
8386 private final WorkflowClientOptions workflowClientOptions ;
@@ -90,6 +93,7 @@ public class TestWorkflowExtension
9093 private final long initialTimeMillis ;
9194
9295 private final Set <Class <?>> supportedParameterTypes = new HashSet <>();
96+ private boolean includesDynamicWorkflow ;
9397
9498 private TestWorkflowExtension (Builder builder ) {
9599 workerOptions = builder .workerOptions ;
@@ -113,6 +117,10 @@ private TestWorkflowExtension(Builder builder) {
113117 supportedParameterTypes .add (Worker .class );
114118
115119 for (Class <?> workflowType : workflowTypes ) {
120+ if (DynamicWorkflow .class .isAssignableFrom (workflowType )) {
121+ includesDynamicWorkflow = true ;
122+ continue ;
123+ }
116124 POJOWorkflowImplMetadata metadata = POJOWorkflowImplMetadata .newInstance (workflowType );
117125 for (POJOWorkflowInterfaceMetadata workflowInterface : metadata .getWorkflowInterfaces ()) {
118126 supportedParameterTypes .add (workflowInterface .getInterfaceClass ());
@@ -129,50 +137,68 @@ public boolean supportsParameter(
129137 ParameterContext parameterContext , ExtensionContext extensionContext )
130138 throws ParameterResolutionException {
131139
132- if (parameterContext .getParameter ().getDeclaringExecutable () instanceof Constructor ) {
140+ Parameter parameter = parameterContext .getParameter ();
141+ if (parameter .getDeclaringExecutable () instanceof Constructor ) {
133142 // Constructor injection is not supported
134143 return false ;
135144 }
136145
137- Class <?> parameterType = parameterContext .getParameter ().getType ();
138- return supportedParameterTypes .contains (parameterType );
146+ Class <?> parameterType = parameter .getType ();
147+ if (supportedParameterTypes .contains (parameterType )) {
148+ return true ;
149+ }
150+
151+ if (!includesDynamicWorkflow ) {
152+ // If no DynamicWorkflow implementation was registered then supportedParameterTypes are the
153+ // only ones types that can be injected
154+ return false ;
155+ }
156+
157+ try {
158+ // If POJOWorkflowInterfaceMetadata can be instantiated then parameterType is a proper
159+ // workflow interface and can be injected
160+ POJOWorkflowInterfaceMetadata .newInstance (parameterType );
161+ return true ;
162+ } catch (Exception e ) {
163+ return false ;
164+ }
139165 }
140166
141167 @ Override
142168 public Object resolveParameter (
143169 ParameterContext parameterContext , ExtensionContext extensionContext )
144170 throws ParameterResolutionException {
145171
146- TestWorkflowEnvironment testEnvironment = getTestEnvironment (extensionContext );
147-
148172 Class <?> parameterType = parameterContext .getParameter ().getType ();
149173 if (parameterType == TestWorkflowEnvironment .class ) {
150- return testEnvironment ;
174+ return getTestEnvironment ( extensionContext ) ;
151175 } else if (parameterType == WorkflowClient .class ) {
152- return testEnvironment .getWorkflowClient ();
176+ return getTestEnvironment ( extensionContext ) .getWorkflowClient ();
153177 } else if (parameterType == WorkflowOptions .class ) {
154- String taskQueue = getTaskQueue (extensionContext );
155- return WorkflowOptions .newBuilder ().setTaskQueue (taskQueue ).build ();
178+ return getWorkflowOptions (extensionContext );
156179 } else if (parameterType == Worker .class ) {
157180 return getWorker (extensionContext );
158181 } else {
159182 // Workflow stub
160- String taskQueue = getTaskQueue (extensionContext );
161- WorkflowOptions workflowOptions =
162- WorkflowOptions .newBuilder ().setTaskQueue (taskQueue ).build ();
163- return testEnvironment .getWorkflowClient ().newWorkflowStub (parameterType , workflowOptions );
183+ return getTestEnvironment (extensionContext )
184+ .getWorkflowClient ()
185+ .newWorkflowStub (parameterType , getWorkflowOptions (extensionContext ));
164186 }
165187 }
166188
167189 @ Override
168190 public void beforeEach (ExtensionContext context ) throws Exception {
191+ long currentInitialTimeMillis =
192+ AnnotationSupport .findAnnotation (context .getElement (), WorkflowInitialTime .class )
193+ .map (annotation -> Instant .parse (annotation .value ()).toEpochMilli ())
194+ .orElse (initialTimeMillis );
169195 TestEnvironmentOptions testOptions =
170196 TestEnvironmentOptions .newBuilder ()
171197 .setWorkflowClientOptions (workflowClientOptions )
172198 .setWorkerFactoryOptions (workerFactoryOptions )
173199 .setUseExternalService (useExternalService )
174200 .setTarget (target )
175- .setInitialTimeMillis (initialTimeMillis )
201+ .setInitialTimeMillis (currentInitialTimeMillis )
176202 .build ();
177203 TestWorkflowEnvironment testEnvironment = TestWorkflowEnvironment .newInstance (testOptions );
178204 String taskQueue =
@@ -187,7 +213,7 @@ public void beforeEach(ExtensionContext context) throws Exception {
187213
188214 setTestEnvironment (context , testEnvironment );
189215 setWorker (context , worker );
190- setTaskQueue (context , taskQueue );
216+ setWorkflowOptions (context , WorkflowOptions . newBuilder (). setTaskQueue ( taskQueue ). build () );
191217 }
192218
193219 @ Override
@@ -219,12 +245,12 @@ private void setWorker(ExtensionContext context, Worker worker) {
219245 getStore (context ).put (WORKER_KEY , worker );
220246 }
221247
222- private String getTaskQueue (ExtensionContext context ) {
223- return getStore (context ).get (TASK_QUEUE_KEY , String .class );
248+ private WorkflowOptions getWorkflowOptions (ExtensionContext context ) {
249+ return getStore (context ).get (WORKFLOW_OPTIONS_KEY , WorkflowOptions .class );
224250 }
225251
226- private void setTaskQueue (ExtensionContext context , String taskQueue ) {
227- getStore (context ).put (TASK_QUEUE_KEY , taskQueue );
252+ private void setWorkflowOptions (ExtensionContext context , WorkflowOptions taskQueue ) {
253+ getStore (context ).put (WORKFLOW_OPTIONS_KEY , taskQueue );
228254 }
229255
230256 private ExtensionContext .Store getStore (ExtensionContext context ) {
0 commit comments