|
@@ -19,8 +19,8 @@ package org.apache.dolphinscheduler.plugin.task.jupyter;
|
|
|
|
|
|
|
|
|
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext;
|
|
|
+import org.apache.dolphinscheduler.spi.utils.DateUtils;
|
|
|
import org.apache.dolphinscheduler.spi.utils.JSONUtils;
|
|
|
-
|
|
|
import org.apache.dolphinscheduler.spi.utils.PropertyUtils;
|
|
|
import org.junit.Assert;
|
|
|
import org.junit.Test;
|
|
@@ -30,16 +30,15 @@ import org.powermock.core.classloader.annotations.PowerMockIgnore;
|
|
|
import org.powermock.core.classloader.annotations.PrepareForTest;
|
|
|
import org.powermock.core.classloader.annotations.SuppressStaticInitializationFor;
|
|
|
import org.powermock.modules.junit4.PowerMockRunner;
|
|
|
-import org.apache.dolphinscheduler.plugin.task.api.TaskConstants;
|
|
|
-
|
|
|
import static org.mockito.ArgumentMatchers.any;
|
|
|
import static org.powermock.api.mockito.PowerMockito.spy;
|
|
|
import static org.powermock.api.mockito.PowerMockito.when;
|
|
|
|
|
|
@RunWith(PowerMockRunner.class)
|
|
|
@PrepareForTest({
|
|
|
- JSONUtils.class,
|
|
|
- PropertyUtils.class,
|
|
|
+ JSONUtils.class,
|
|
|
+ PropertyUtils.class,
|
|
|
+ DateUtils.class
|
|
|
})
|
|
|
@PowerMockIgnore({"javax.*"})
|
|
|
@SuppressStaticInitializationFor("org.apache.dolphinscheduler.spi.utils.PropertyUtils")
|
|
@@ -99,6 +98,39 @@ public class JupyterTaskTest {
|
|
|
"--progress-bar");
|
|
|
}
|
|
|
|
|
|
+ @Test
|
|
|
+ public void testBuildJupyterCommandWithRequirements() throws Exception {
|
|
|
+ String parameters = buildJupyterCommandWithRequirements();
|
|
|
+ TaskExecutionContext taskExecutionContext = PowerMockito.mock(TaskExecutionContext.class);
|
|
|
+ when(taskExecutionContext.getTaskParams()).thenReturn(parameters);
|
|
|
+ PowerMockito.mockStatic(PropertyUtils.class);
|
|
|
+ when(PropertyUtils.getString(any())).thenReturn("/opt/anaconda3/etc/profile.d/conda.sh");
|
|
|
+ PowerMockito.mockStatic(DateUtils.class);
|
|
|
+ when(DateUtils.getTimestampString()).thenReturn("123456789");
|
|
|
+ JupyterTask jupyterTask = spy(new JupyterTask(taskExecutionContext));
|
|
|
+ jupyterTask.init();
|
|
|
+ Assert.assertEquals(jupyterTask.buildCommand(),
|
|
|
+ "set +e \n " +
|
|
|
+ "source /opt/anaconda3/etc/profile.d/conda.sh && " +
|
|
|
+ "conda create -n jupyter-tmp-env-123456789 -y && " +
|
|
|
+ "conda activate jupyter-tmp-env-123456789 && " +
|
|
|
+ "pip install -r requirements.txt && " +
|
|
|
+ "papermill " +
|
|
|
+ "/test/input_note.ipynb " +
|
|
|
+ "/test/output_note.ipynb " +
|
|
|
+ "--parameters city Shanghai " +
|
|
|
+ "--parameters factor 0.01 " +
|
|
|
+ "--kernel python3 " +
|
|
|
+ "--engine default_engine " +
|
|
|
+ "--execution-timeout 10 " +
|
|
|
+ "--start-timeout 3 " +
|
|
|
+ "--version " +
|
|
|
+ "--inject-paths " +
|
|
|
+ "--progress-bar \n " +
|
|
|
+ "conda deactivate && conda remove --name jupyter-tmp-env-123456789 --all -y"
|
|
|
+ );
|
|
|
+ }
|
|
|
+
|
|
|
private String buildJupyterCommandWithLocalEnv() {
|
|
|
JupyterParameters jupyterParameters = new JupyterParameters();
|
|
|
jupyterParameters.setCondaEnvName("jupyter-lab");
|
|
@@ -127,4 +159,18 @@ public class JupyterTaskTest {
|
|
|
return JSONUtils.toJsonString(jupyterParameters);
|
|
|
}
|
|
|
|
|
|
+ private String buildJupyterCommandWithRequirements() {
|
|
|
+ JupyterParameters jupyterParameters = new JupyterParameters();
|
|
|
+ jupyterParameters.setCondaEnvName("requirements.txt");
|
|
|
+ jupyterParameters.setInputNotePath("/test/input_note.ipynb");
|
|
|
+ jupyterParameters.setOutputNotePath("/test/output_note.ipynb");
|
|
|
+ jupyterParameters.setParameters("{\"city\": \"Shanghai\", \"factor\": \"0.01\"}");
|
|
|
+ jupyterParameters.setKernel("python3");
|
|
|
+ jupyterParameters.setEngine("default_engine");
|
|
|
+ jupyterParameters.setExecutionTimeout("10");
|
|
|
+ jupyterParameters.setStartTimeout("3");
|
|
|
+ jupyterParameters.setOthers("--version");
|
|
|
+ return JSONUtils.toJsonString(jupyterParameters);
|
|
|
+ }
|
|
|
+
|
|
|
}
|