FlyteAdminClientRunner implements FlyteRunner { @Override public void poll(FlyteExecutionId flyteExecutionId, RunState runState) throws PollingException { requireNonNull(flyteExecutionId, "flyteExecutionId"); requireNonNull(runState, "runState"); try { final ExecutionOuterClass.Execution execution = flyteAdminClient.getExecution(flyteExecutionId.project(), flyteExecutionId.domain(), flyteExecutionId.name()); emitFlyteEvents(execution, runState); } catch (StatusRuntimeException e) { if (e.getStatus().getCode() == Status.Code.NOT_FOUND) { throw new ExecutionNotFoundException(flyteExecutionId, e); } throw new PollingException(flyteExecutionId, e); } catch (Exception e) { throw new PollingException(flyteExecutionId, e); } } @VisibleForTesting  FlyteAdminClientRunner(final String runnerId,
                         final FlyteAdminClient flyteAdminClient,
                         final StateManager stateManager,
                         final Duration terminateDanglingFlyteExecInterval,
                         final ScheduledExecutorService scheduledExecutor,
                         final Time time);   FlyteAdminClientRunner(final String runnerId,
                         final FlyteAdminClient flyteAdminClient,
                         final StateManager stateManager); @Override String createExecution(final RunState runState, final String execName, final FlyteExecConf flyteExecConf,
                                final Map<String, String> annotations); @Override void terminateExecution(RunState runState, FlyteExecutionId flyteExecutionId); @Override void poll(FlyteExecutionId flyteExecutionId, RunState runState); @Override void close();  }
@Test public void testPollFlyteExecutionIdCannotBeNull() { assertThrows( NullPointerException.class, () -> flyteRunner.poll(null, RUN_STATE) ); }
@Test public void testPollRunStateCannotBeNull() { assertThrows( NullPointerException.class, () -> flyteRunner.poll(FLYTE_EXECUTION_ID, null) ); }
@Test public void testTransitionRunningToTerminateSuccessfulRun() throws Exception { when(flyteAdminClient.getExecution("flyte-test", "testing", "execution-name")).thenReturn( ExecutionOuterClass.Execution .newBuilder() .setClosure(ExecutionOuterClass.ExecutionClosure.newBuilder() .setPhase(Execution.WorkflowExecution.Phase.SUCCEEDED).build()) .build()); flyteRunner.poll(FLYTE_EXECUTION_ID, RUN_STATE); verify(stateManager, timeout(60_000)).receive(Event.terminate(WORKFLOW_INSTANCE, Optional.of(0))); }
@Test @Parameters(method = "parametersForTestTransitionPollingToTerminateExitCodes") public void testTransitionRunningToTerminateExitCodes(Execution.WorkflowExecution.Phase phase, String flyteExitCode, int exitCode) throws Exception { when(flyteAdminClient.getExecution("flyte-test", "testing", "execution-name")).thenReturn( ExecutionOuterClass.Execution .newBuilder() .setClosure(ExecutionOuterClass.ExecutionClosure.newBuilder().setError( Execution.ExecutionError .newBuilder() .setCode(flyteExitCode) .build()) .setPhase(phase).build()) .build()); flyteRunner.poll(FLYTE_EXECUTION_ID, RUN_STATE); verify(stateManager, timeout(60_000)).receive(Event.terminate(WORKFLOW_INSTANCE, Optional.of(exitCode))); }
@Test @Parameters(method = "parametersForTestTransitionPollingToTerminateExitCodes") public void testTransitionSubmittedToTerminateExitCodes(Execution.WorkflowExecution.Phase phase, String flyteExitCode, int exitCode) throws Exception { when(flyteAdminClient.getExecution("flyte-test", "testing", "execution-name")).thenReturn( ExecutionOuterClass.Execution .newBuilder() .setClosure(ExecutionOuterClass.ExecutionClosure.newBuilder().setError( Execution.ExecutionError .newBuilder() .setCode(flyteExitCode) .build()) .setPhase(phase).build()) .build()); flyteRunner.poll(FLYTE_EXECUTION_ID, RUN_STATE_SUBMITTED); verify(stateManager, timeout(60_000)).receive(Event.started(WORKFLOW_INSTANCE)); verify(stateManager, timeout(60_000)).receive(Event.terminate(WORKFLOW_INSTANCE, Optional.of(exitCode))); }
@Test public void testTransitionSubmittedToRunning() throws Exception { when(flyteAdminClient.getExecution("flyte-test", "testing", "execution-name")).thenReturn( ExecutionOuterClass.Execution .newBuilder() .setClosure(ExecutionOuterClass.ExecutionClosure.newBuilder() .setPhase(Execution.WorkflowExecution.Phase.RUNNING).build()) .build()); flyteRunner.poll(FLYTE_EXECUTION_ID, RUN_STATE_SUBMITTED); verify(stateManager, timeout(60_000)).receive(Event.started(WORKFLOW_INSTANCE)); }
@Test public void testTransitionSubmittedToTerminate() throws Exception { when(flyteAdminClient.getExecution("flyte-test", "testing", "execution-name")).thenReturn( ExecutionOuterClass.Execution .newBuilder() .setClosure(ExecutionOuterClass.ExecutionClosure.newBuilder() .setPhase(Execution.WorkflowExecution.Phase.SUCCEEDED).build()) .build()); flyteRunner.poll(FLYTE_EXECUTION_ID, RUN_STATE_SUBMITTED); verify(stateManager, timeout(60_000)).receive(Event.started(WORKFLOW_INSTANCE)); verify(stateManager, timeout(60_000)).receive(Event.terminate(WORKFLOW_INSTANCE, Optional.of(SUCCESS_EXIT_CODE))); }
@Test public void testTransitionRunningToRunningNoEventEmit() throws Exception { when(flyteAdminClient.getExecution("flyte-test", "testing", "execution-name")).thenReturn( ExecutionOuterClass.Execution .newBuilder() .setClosure(ExecutionOuterClass.ExecutionClosure.newBuilder() .setPhase(Execution.WorkflowExecution.Phase.RUNNING).build()) .build()); RunState state = runState(RunState.State.RUNNING); flyteRunner.poll(FLYTE_EXECUTION_ID, state); verify(stateManager, never()).receive(Event.started(WORKFLOW_INSTANCE)); }
@Test public void testPollingExceptionFlyteAdminClientExecutionNotFound() { doThrow(new StatusRuntimeException(Status.NOT_FOUND)) .when(flyteAdminClient).getExecution(anyString(), anyString(), anyString()); assertThrows( FlyteRunner.ExecutionNotFoundException.class, () -> flyteRunner.poll(FLYTE_EXECUTION_ID, RUN_STATE)); }
@Test public void testPollingExceptionFlyteAdminClientExecution() { doThrow(new StatusRuntimeException(Status.INTERNAL)) .when(flyteAdminClient).getExecution(anyString(), anyString(), anyString()); assertThrows( FlyteRunner.PollingException.class, () -> flyteRunner.poll(FLYTE_EXECUTION_ID, RUN_STATE)); }
@Test public void testPollingException() { doThrow(new RuntimeException()) .when(flyteAdminClient).getExecution(anyString(), anyString(), anyString()); assertThrows( FlyteRunner.PollingException.class, () -> flyteRunner.poll(FLYTE_EXECUTION_ID, RUN_STATE)); }
@Test public void testUndefinedShouldNotInteractWithStateManager() throws Exception { when(flyteAdminClient.getExecution("flyte-test", "testing", "execution-name")).thenReturn( ExecutionOuterClass.Execution .newBuilder() .setClosure(ExecutionOuterClass.ExecutionClosure.newBuilder() .setPhase(Execution.WorkflowExecution.Phase.UNDEFINED).build()) .build()); flyteRunner.poll(FLYTE_EXECUTION_ID, RUN_STATE); verifyNoInteractions(stateManager); }