diff --git a/dolphinscheduler-task-plugin/dolphinscheduler-task-openmldb/src/test/java/org/apache/dolphinscheduler/plugin/task/openmldb/OpenmldbTaskTest.java b/dolphinscheduler-task-plugin/dolphinscheduler-task-openmldb/src/test/java/org/apache/dolphinscheduler/plugin/task/openmldb/OpenmldbTaskTest.java
index ac17696c63..c608fa8fb3 100644
--- a/dolphinscheduler-task-plugin/dolphinscheduler-task-openmldb/src/test/java/org/apache/dolphinscheduler/plugin/task/openmldb/OpenmldbTaskTest.java
+++ b/dolphinscheduler-task-plugin/dolphinscheduler-task-openmldb/src/test/java/org/apache/dolphinscheduler/plugin/task/openmldb/OpenmldbTaskTest.java
@@ -25,9 +25,9 @@ import org.apache.dolphinscheduler.plugin.task.api.parameters.AbstractParameters
 import java.util.HashMap;
 import java.util.Map;
 
-import org.junit.Assert;
-import org.junit.Test;
-import org.powermock.reflect.Whitebox;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
 
 public class OpenmldbTaskTest {
     static class MockOpenmldbTask extends OpenmldbTask {
@@ -55,36 +55,39 @@ public class OpenmldbTaskTest {
         OpenmldbTask openmldbTask = createOpenmldbTask();
         String pythonFile = "test.py";
         String result1 = openmldbTask.buildPythonExecuteCommand(pythonFile);
-        Assert.assertEquals("python3 test.py", result1);
+        Assertions.assertEquals("python3 test.py", result1);
     }
 
     @Test
     public void buildSQLWithComment() throws Exception {
-        OpenmldbTask openmldbTask = createOpenmldbTask();
+        TaskExecutionContext taskExecutionContext = Mockito.mock(TaskExecutionContext.class);
         OpenmldbParameters openmldbParameters = new OpenmldbParameters();
         openmldbParameters.setExecuteMode("offline");
+        openmldbParameters.setZk("localhost:2181");
+        openmldbParameters.setZkPath("dolphinscheduler");
         String rawSQLScript = "select * from users\r\n"
                 + "-- some comment\n"
                 + "inner join order on users.order_id = order.id; \n\n;"
                 + "select * from users;";
         openmldbParameters.setSql(rawSQLScript);
-        Whitebox.setInternalState(openmldbTask, "openmldbParameters", openmldbParameters);
+        Mockito.when(taskExecutionContext.getTaskParams()).thenReturn(JSONUtils.toJsonString(openmldbParameters));
+        OpenmldbTask openmldbTask = new OpenmldbTask(taskExecutionContext);
+        openmldbTask.init();
         OpenmldbParameters internal = (OpenmldbParameters) openmldbTask.getParameters();
-        Assert.assertNotNull(internal);
-        Assert.assertEquals(internal.getExecuteMode(), "offline");
+        Assertions.assertNotNull(internal);
+        Assertions.assertEquals(internal.getExecuteMode(), "offline");
 
         String result1 = openmldbTask.buildPythonScriptContent();
-        Assert.assertEquals("import openmldb\n"
-                        + "import sqlalchemy as db\n"
-                        + "engine = db.create_engine('openmldb:///?zk=null&zkPath=null')\n"
-                        + "con = engine.connect()\n"
-                        + "con.execute(\"set @@execute_mode='offline';\")\n"
-                        + "con.execute(\"set @@sync_job=true\")\n"
-                        + "con.execute(\"set @@job_timeout=1800000\")\n"
-                        + "con.execute(\"select * from users\\n-- some comment\\ninner join order on users.order_id = "
-                        + "order.id\")\n"
-                        + "con.execute(\"select * from users\")\n"
-                , result1);
+        Assertions.assertEquals("import openmldb\n"
+                + "import sqlalchemy as db\n"
+                + "engine = db.create_engine('openmldb:///?zk=localhost:2181&zkPath=dolphinscheduler')\n"
+                + "con = engine.connect()\n"
+                + "con.execute(\"set @@execute_mode='offline';\")\n"
+                + "con.execute(\"set @@sync_job=true\")\n"
+                + "con.execute(\"set @@job_timeout=1800000\")\n"
+                + "con.execute(\"select * from users\\n-- some comment\\ninner join order on users.order_id = "
+                + "order.id\")\n"
+                + "con.execute(\"select * from users\")\n", result1);
     }
 
 }
\ No newline at end of file