|
@@ -29,6 +29,7 @@ import static org.powermock.api.mockito.PowerMockito.when;
|
|
|
|
|
|
import com.fasterxml.jackson.databind.ObjectMapper;
|
|
|
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext;
|
|
|
+import org.apache.dolphinscheduler.spi.utils.DateUtils;
|
|
|
import org.apache.dolphinscheduler.spi.utils.JSONUtils;
|
|
|
|
|
|
|
|
@@ -51,9 +52,10 @@ import java.util.Map;
|
|
|
|
|
|
@RunWith(PowerMockRunner.class)
|
|
|
@PrepareForTest({
|
|
|
- ZeppelinTask.class,
|
|
|
- ZeppelinClient.class,
|
|
|
- ObjectMapper.class,
|
|
|
+ ZeppelinTask.class,
|
|
|
+ ZeppelinClient.class,
|
|
|
+ ObjectMapper.class,
|
|
|
+ DateUtils.class
|
|
|
})
|
|
|
@PowerMockIgnore({"javax.*"})
|
|
|
public class ZeppelinTaskTest {
|
|
@@ -62,6 +64,8 @@ public class ZeppelinTaskTest {
|
|
|
private static final String MOCK_PARAGRAPH_ID = "paragraph_1648793472526_1771221396";
|
|
|
private static final String MOCK_PARAMETERS = "{\"key1\": \"value1\", \"key2\": \"value2\"}";
|
|
|
private static final String MOCK_REST_ENDPOINT = "localhost:8080";
|
|
|
+ private static final String MOCK_CLONE_NOTE_ID = "3GYJR92R8";
|
|
|
+ private static final String MOCK_PRODUCTION_DIRECTORY = "/prod/";
|
|
|
private final ObjectMapper mapper = new ObjectMapper();
|
|
|
|
|
|
private ZeppelinClient zClient;
|
|
@@ -161,6 +165,37 @@ public class ZeppelinTaskTest {
|
|
|
Assert.assertEquals(EXIT_CODE_SUCCESS, this.zeppelinTask.getExitStatusCode());
|
|
|
}
|
|
|
|
|
|
+ @Test
|
|
|
+ public void testHandleWithNoteExecutionSuccessWithProductionSetting() throws Exception {
|
|
|
+ String zeppelinParametersWithNoParagraphId = buildZeppelinTaskParametersWithProductionSetting();
|
|
|
+ TaskExecutionContext taskExecutionContext = PowerMockito.mock(TaskExecutionContext.class);
|
|
|
+ PowerMockito.mockStatic(DateUtils.class);
|
|
|
+ when(taskExecutionContext.getTaskParams()).thenReturn(zeppelinParametersWithNoParagraphId);
|
|
|
+ this.zeppelinTask = spy(new ZeppelinTask(taskExecutionContext));
|
|
|
+
|
|
|
+ // mock zClient and note result
|
|
|
+ this.zClient = mock(ZeppelinClient.class);
|
|
|
+ this.noteResult = mock(NoteResult.class);
|
|
|
+
|
|
|
+ // use mocked zClient in zeppelinTask
|
|
|
+ doReturn(this.zClient).when(this.zeppelinTask, "getZeppelinClient");
|
|
|
+ when(this.zClient.cloneNote(any(String.class), any(String.class))).thenReturn(MOCK_CLONE_NOTE_ID);
|
|
|
+ when(this.zClient.executeNote(any(), any(Map.class))).thenReturn(this.noteResult);
|
|
|
+ when(paragraphResult.getResultInText()).thenReturn("mock-zeppelin-paragraph-execution-result");
|
|
|
+ this.zeppelinTask.init();
|
|
|
+ when(this.paragraphResult.getStatus()).thenReturn(Status.FINISHED);
|
|
|
+ when(DateUtils.getTimestampString()).thenReturn("123456789");
|
|
|
+ this.zeppelinTask.handle();
|
|
|
+ Mockito.verify(this.zClient).cloneNote(
|
|
|
+ MOCK_NOTE_ID,
|
|
|
+ String.format("%s%s_%s", MOCK_PRODUCTION_DIRECTORY, MOCK_NOTE_ID, "123456789"));
|
|
|
+ Mockito.verify(this.zClient).executeNote(MOCK_CLONE_NOTE_ID,
|
|
|
+ (Map<String, String>) mapper.readValue(MOCK_PARAMETERS, Map.class));
|
|
|
+ Mockito.verify(this.noteResult).getParagraphResultList();
|
|
|
+ Mockito.verify(this.zClient).deleteNote(MOCK_CLONE_NOTE_ID);
|
|
|
+ Assert.assertEquals(EXIT_CODE_SUCCESS, this.zeppelinTask.getExitStatusCode());
|
|
|
+ }
|
|
|
+
|
|
|
private String buildZeppelinTaskParameters() {
|
|
|
ZeppelinParameters zeppelinParameters = new ZeppelinParameters();
|
|
|
zeppelinParameters.setNoteId(MOCK_NOTE_ID);
|
|
@@ -179,4 +214,15 @@ public class ZeppelinTaskTest {
|
|
|
|
|
|
return JSONUtils.toJsonString(zeppelinParameters);
|
|
|
}
|
|
|
+
|
|
|
+ private String buildZeppelinTaskParametersWithProductionSetting() {
|
|
|
+ ZeppelinParameters zeppelinParameters = new ZeppelinParameters();
|
|
|
+ zeppelinParameters.setNoteId(MOCK_NOTE_ID);
|
|
|
+ zeppelinParameters.setParameters(MOCK_PARAMETERS);
|
|
|
+ zeppelinParameters.setRestEndpoint(MOCK_REST_ENDPOINT);
|
|
|
+ zeppelinParameters.setProductionNoteDirectory(MOCK_PRODUCTION_DIRECTORY);
|
|
|
+
|
|
|
+ return JSONUtils.toJsonString(zeppelinParameters);
|
|
|
+ }
|
|
|
+
|
|
|
}
|