|
|
|
@ -17,121 +17,112 @@
|
|
|
|
|
|
|
|
|
|
package org.apache.dolphinscheduler.plugin.task.jupyter; |
|
|
|
|
|
|
|
|
|
import static org.mockito.Mockito.doReturn; |
|
|
|
|
import static org.mockito.Mockito.spy; |
|
|
|
|
import static org.mockito.Mockito.when; |
|
|
|
|
|
|
|
|
|
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; |
|
|
|
|
import org.junit.runner.RunWith; |
|
|
|
|
import org.powermock.api.mockito.PowerMockito; |
|
|
|
|
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 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, |
|
|
|
|
DateUtils.class |
|
|
|
|
}) |
|
|
|
|
@PowerMockIgnore({"javax.*"}) |
|
|
|
|
@SuppressStaticInitializationFor("org.apache.dolphinscheduler.spi.utils.PropertyUtils") |
|
|
|
|
import org.mockito.Mockito; |
|
|
|
|
import org.mockito.junit.MockitoJUnitRunner; |
|
|
|
|
|
|
|
|
|
@RunWith(MockitoJUnitRunner.class) |
|
|
|
|
public class JupyterTaskTest { |
|
|
|
|
|
|
|
|
|
private static final String EXPECTED_JUPYTER_TASK_COMMAND_USE_LOCAL_CONDA_ENV = |
|
|
|
|
"source /opt/anaconda3/etc/profile.d/conda.sh && " + |
|
|
|
|
"conda activate jupyter-lab && " + |
|
|
|
|
"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"; |
|
|
|
|
|
|
|
|
|
private static final String EXPECTED_JUPYTER_TASK_COMMAND_USE_PACKED_CONDA_ENV = |
|
|
|
|
"source /opt/anaconda3/etc/profile.d/conda.sh && " + |
|
|
|
|
"mkdir jupyter_env && " + |
|
|
|
|
"tar -xzf jupyter.tar.gz -C jupyter_env && " + |
|
|
|
|
"source jupyter_env/bin/activate && " + |
|
|
|
|
"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"; |
|
|
|
|
|
|
|
|
|
private static final String EXPECTED_JUPYTER_TASK_COMMAND_USE_PIP_REQUIREMENTS = |
|
|
|
|
"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"; |
|
|
|
|
|
|
|
|
|
@Test |
|
|
|
|
public void testBuildJupyterCommandWithLocalEnv() throws Exception { |
|
|
|
|
String parameters = buildJupyterCommandWithLocalEnv(); |
|
|
|
|
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"); |
|
|
|
|
JupyterTask jupyterTask = spy(new JupyterTask(taskExecutionContext)); |
|
|
|
|
public void jupyterTaskUseLocalCondaEnv() throws Exception { |
|
|
|
|
String jupyterTaskParameters = buildJupyterTaskUseLocalCondaEnvCommand(); |
|
|
|
|
JupyterTask jupyterTask = prepareJupyterTaskForTest(jupyterTaskParameters); |
|
|
|
|
jupyterTask.init(); |
|
|
|
|
Assert.assertEquals(jupyterTask.buildCommand(), |
|
|
|
|
"source /opt/anaconda3/etc/profile.d/conda.sh && " + |
|
|
|
|
"conda activate jupyter-lab && " + |
|
|
|
|
"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"); |
|
|
|
|
Assert.assertEquals(jupyterTask.buildCommand(), EXPECTED_JUPYTER_TASK_COMMAND_USE_LOCAL_CONDA_ENV); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
@Test |
|
|
|
|
public void testBuildJupyterCommandWithPackedEnv() throws Exception { |
|
|
|
|
String parameters = buildJupyterCommandWithPackedEnv(); |
|
|
|
|
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"); |
|
|
|
|
JupyterTask jupyterTask = spy(new JupyterTask(taskExecutionContext)); |
|
|
|
|
public void jupyterTaskUsePackedCondaEnv() throws Exception { |
|
|
|
|
String jupyterTaskParameters = buildJupyterTaskUsePackedCondaEnvCommand(); |
|
|
|
|
JupyterTask jupyterTask = prepareJupyterTaskForTest(jupyterTaskParameters); |
|
|
|
|
jupyterTask.init(); |
|
|
|
|
Assert.assertEquals(jupyterTask.buildCommand(), |
|
|
|
|
"source /opt/anaconda3/etc/profile.d/conda.sh && " + |
|
|
|
|
"mkdir jupyter_env && " + |
|
|
|
|
"tar -xzf jupyter.tar.gz -C jupyter_env && " + |
|
|
|
|
"source jupyter_env/bin/activate && " + |
|
|
|
|
"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"); |
|
|
|
|
Assert.assertEquals(jupyterTask.buildCommand(), EXPECTED_JUPYTER_TASK_COMMAND_USE_PACKED_CONDA_ENV); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
@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); |
|
|
|
|
public void jupyterTaskUsePipRequirements() throws Exception { |
|
|
|
|
String jupyterTaskParameters = buildJupyterTaskUsePipRequirementsCommand(); |
|
|
|
|
JupyterTask jupyterTask = prepareJupyterTaskForTest(jupyterTaskParameters); |
|
|
|
|
Mockito.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" |
|
|
|
|
); |
|
|
|
|
Assert.assertEquals(jupyterTask.buildCommand(), EXPECTED_JUPYTER_TASK_COMMAND_USE_PIP_REQUIREMENTS); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
private JupyterTask prepareJupyterTaskForTest(final String jupyterTaskParameters) { |
|
|
|
|
TaskExecutionContext taskExecutionContext = Mockito.mock(TaskExecutionContext.class); |
|
|
|
|
when(taskExecutionContext.getTaskParams()).thenReturn(jupyterTaskParameters); |
|
|
|
|
JupyterTask jupyterTask = spy(new JupyterTask(taskExecutionContext)); |
|
|
|
|
doReturn("/opt/anaconda3/etc/profile.d/conda.sh").when(jupyterTask).readCondaPath(); |
|
|
|
|
return jupyterTask; |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
private String buildJupyterCommandWithLocalEnv() { |
|
|
|
|
private String buildJupyterTaskUseLocalCondaEnvCommand() { |
|
|
|
|
JupyterParameters jupyterParameters = new JupyterParameters(); |
|
|
|
|
jupyterParameters.setCondaEnvName("jupyter-lab"); |
|
|
|
|
jupyterParameters.setInputNotePath("/test/input_note.ipynb"); |
|
|
|
@ -145,7 +136,7 @@ public class JupyterTaskTest {
|
|
|
|
|
return JSONUtils.toJsonString(jupyterParameters); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
private String buildJupyterCommandWithPackedEnv() { |
|
|
|
|
private String buildJupyterTaskUsePackedCondaEnvCommand() { |
|
|
|
|
JupyterParameters jupyterParameters = new JupyterParameters(); |
|
|
|
|
jupyterParameters.setCondaEnvName("jupyter.tar.gz"); |
|
|
|
|
jupyterParameters.setInputNotePath("/test/input_note.ipynb"); |
|
|
|
@ -159,7 +150,7 @@ public class JupyterTaskTest {
|
|
|
|
|
return JSONUtils.toJsonString(jupyterParameters); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
private String buildJupyterCommandWithRequirements() { |
|
|
|
|
private String buildJupyterTaskUsePipRequirementsCommand() { |
|
|
|
|
JupyterParameters jupyterParameters = new JupyterParameters(); |
|
|
|
|
jupyterParameters.setCondaEnvName("requirements.txt"); |
|
|
|
|
jupyterParameters.setInputNotePath("/test/input_note.ipynb"); |
|
|
|
|