diff --git a/src/main/java/com/liveramp/daemon_lib/utils/JobletProcessHandler.java b/src/main/java/com/liveramp/daemon_lib/utils/JobletProcessHandler.java index 549ef66..9b73e32 100644 --- a/src/main/java/com/liveramp/daemon_lib/utils/JobletProcessHandler.java +++ b/src/main/java/com/liveramp/daemon_lib/utils/JobletProcessHandler.java @@ -41,8 +41,19 @@ public void onRemove(ProcessDefinition watchedProcess) throws DaemonExce } if (jobletStatusManager.exists(identifier)) { + JobletStatus status; + try { + status = jobletStatusManager.getStatus(identifier); + } catch (Exception e) { + try { + failureCallback.callback(jobletConfig); + } catch (Exception callbackException) { + throw new DaemonException("Failure callback triggered due to an exception in jobletStatusManager. Failure callback, however, also failed.", callbackException); + } + throw e; + } + try { - JobletStatus status = jobletStatusManager.getStatus(identifier); switch (status) { case DONE: LOG.info("Process succeeded - PID: " + watchedProcess.getPid()); diff --git a/src/test/java/com/liveramp/daemon_lib/utils/TestJobletProcessHandler.java b/src/test/java/com/liveramp/daemon_lib/utils/TestJobletProcessHandler.java new file mode 100644 index 0000000..1220d5b --- /dev/null +++ b/src/test/java/com/liveramp/daemon_lib/utils/TestJobletProcessHandler.java @@ -0,0 +1,103 @@ +package com.liveramp.daemon_lib.utils; + +import java.io.IOException; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.runners.MockitoJUnitRunner; + +import com.liveramp.daemon_lib.JobletCallback; +import com.liveramp.daemon_lib.JobletConfig; +import com.liveramp.daemon_lib.executors.processes.ProcessDefinition; +import com.liveramp.daemon_lib.executors.processes.ProcessMetadata; +import com.liveramp.daemon_lib.tracking.JobletStatus; +import com.liveramp.daemon_lib.tracking.JobletStatusManager; + +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@SuppressWarnings("unchecked") +@RunWith(MockitoJUnitRunner.class) +public class TestJobletProcessHandler { + private final String IDENTIFIER = "id"; + + @Mock + private JobletCallback successCallback; + + @Mock + private JobletCallback failureCallback; + + @Mock + private JobletConfigStorage jobletConfigStorage; + + @Mock + private JobletStatusManager jobletStatusManager; + + @Mock + private ProcessDefinition processDefinition; + + @Mock + private ProcessMetadata processMetadata; + + @Mock + private JobletConfig jobletConfig; + + private JobletProcessHandler jobletProcessHandler; + + @Before + public void setup() throws IOException, ClassNotFoundException { + when(processDefinition.getMetadata()).thenReturn(processMetadata); + when(processMetadata.getIdentifier()).thenReturn(IDENTIFIER); + + when(jobletConfigStorage.loadConfig(IDENTIFIER)).thenReturn(jobletConfig); + + when(jobletStatusManager.exists(IDENTIFIER)).thenReturn(true); + + jobletProcessHandler = new JobletProcessHandler<>(successCallback, failureCallback, jobletConfigStorage, jobletStatusManager); + } + + + @Test + public void testOnRemoveAttemptsSuccessCallbackWhenStatusIsDone() throws DaemonException { + when(jobletStatusManager.getStatus(IDENTIFIER)).thenReturn(JobletStatus.DONE); + jobletProcessHandler.onRemove(processDefinition); + verify(successCallback, times(1)).callback(jobletConfig); + } + + @Test + public void testOnRemoveAttemptsFailureCallbackWhenStatusIsInProgress() throws DaemonException { + when(jobletStatusManager.getStatus(IDENTIFIER)).thenReturn(JobletStatus.IN_PROGRESS); + jobletProcessHandler.onRemove(processDefinition); + verify(failureCallback, times(1)).callback(jobletConfig); + } + + @Test + public void testOnRemoveAttemptsFailureCallbackOnExceptionThenReThrows() throws DaemonException { + when(jobletStatusManager.getStatus(IDENTIFIER)).thenThrow(new IllegalArgumentException()); + try { + jobletProcessHandler.onRemove(processDefinition); + Assert.fail(); + } catch (IllegalArgumentException ignored) { + + } + verify(failureCallback, times(1)).callback(jobletConfig); + } + + @Test + public void testOnRemoveAttemptsFailureCallbackOnExceptionAndStillFails() throws DaemonException { + when(jobletStatusManager.getStatus(IDENTIFIER)).thenThrow(new IllegalArgumentException()); + doThrow(new RuntimeException()).when(failureCallback).callback(jobletConfig); + try { + jobletProcessHandler.onRemove(processDefinition); + Assert.fail(); + } catch (DaemonException ignored) { + + } + } + +}