RoutingFlyteRunner extends AbstractRoutingRunner<FlyteRunner> implements FlyteRunner { @Override public String createExecution(RunState runState, String name, FlyteExecConf flyteExecConf, Map<String, String> annotations) throws CreateExecutionException { return runner(runState).createExecution(runState, name, flyteExecConf, annotations); } RoutingFlyteRunner(
      FlyteRunnerFactory runnerFactory,
      Function<RunState, String> runnerId); @Override String createExecution(RunState runState, String name, FlyteExecConf flyteExecConf,
                                Map<String, String> annotations); @Override void terminateExecution(RunState runState, FlyteExecutionId flyteExecutionId); @Override void poll(FlyteExecutionId flyteExecutionId, RunState runState);  }
@Test public void shouldCreateRunnerOnCreateExecution() throws FlyteRunner.CreateExecutionException { flyteRunner.createExecution(runState, EXEC_NAME, execConf, ANNOTATIONS); assertThatCreateCountersContains("default"); verify(createdRunners.get("default")).createExecution(runState, EXEC_NAME, execConf, ANNOTATIONS); }
@Test public void testCreatesOnlyOneRunnerPerRunnerId() throws Exception { flyteRunner.createExecution(runState, EXEC_NAME, execConf, ANNOTATIONS); flyteRunner.createExecution(runState, EXEC_NAME, execConf, ANNOTATIONS); assertThatCreateCountersContains("default"); }
@Test public void testSwitchesRunners() throws Exception { when(runnerId.apply(runState)).thenReturn("id-1", "id-2"); flyteRunner.createExecution(runState, EXEC_NAME, execConf, ANNOTATIONS); flyteRunner.createExecution(runState, EXEC_NAME, execConf, ANNOTATIONS); assertThatCreateCountersContains("id-1", "id-2"); }
@Test public void testCreatedRunnersAreClosed() throws Exception { when(runnerId.apply(runState)).thenReturn("id-1", "id-2"); flyteRunner.createExecution(runState, EXEC_NAME, execConf, ANNOTATIONS); flyteRunner.createExecution(runState, EXEC_NAME, execConf, ANNOTATIONS); flyteRunner.close(); assertThatCreateCountersContains("id-1", "id-2"); for (var runner : createdRunners.values()) { verify(runner).close(); } }