fix mock framework

Author: iaorekhov-1980

fe/fe-core/src/test/java/org/apache/doris/nereids/parser/EncryptSQLTest.java

@@ -446,35 +446,65 @@ public void testEncryption() throws Exception {
     @Test
     public void testCreateUserPasswordMasking() throws Exception {
         ctx.setDatabase("test");
-        new MockUp<StmtExecutor>() {
-            @Mock
-            public boolean isForwardToMaster() {
-                return false;
-            }
-        };
-        ctx.setEnv(env);
-        Config.enable_nereids_load = true;
-        // testing for https://github.com/apache/doris/issues/62140
-        String sql = "CREATE USER 'test_user62140'@'%' IDENTIFIED BY '123456'";
-        String res = "CREATE USER 'test_user62140'@'%' IDENTIFIED BY '*XXX'";
-        parseAndCheck(sql, res);
+        try (MockedConstruction<StmtExecutor> ignored = Mockito.mockConstruction(StmtExecutor.class,
+                Mockito.withSettings().defaultAnswer(Mockito.CALLS_REAL_METHODS),
+                (mock, context) -> {
+                    Mockito.doReturn(false).when(mock).isForwardToMaster();
+                    Profile profile = new Profile(false, 0, 0);
+                    Deencapsulation.setField(mock, "profile", profile);
+                    Mockito.doReturn(profile).when(mock).getProfile();
+                    Mockito.doReturn(ctx).when(mock).getContext();
+                    Mockito.doNothing().when(mock).execute();
+                    Deencapsulation.setField(mock, "context", ctx);
+                    if (context.arguments().size() >= 2
+                            && context.arguments().get(1) instanceof StatementBase) {
+                        Deencapsulation.setField(mock, "parsedStmt",
+                                context.arguments().get(1));
+                    }
+                    if (ctx.getStatementContext() == null) {
+                        ctx.setStatementContext(new StatementContext());
+                    }
+                })) {
+            ctx.setEnv(env);
+            ctx.setCurrentUserIdentity(UserIdentity.ROOT);
+            Config.enable_nereids_load = true;
+            // testing for https://github.com/apache/doris/issues/62140
+            String sql = "CREATE USER 'test_user62140'@'%' IDENTIFIED BY '123456'";
+            String res = "CREATE USER 'test_user62140'@'%' IDENTIFIED BY '*XXX'";
+            parseAndCheck(sql, res);
+        }
     }
 
     @Test
     public void testAlterUserPasswordMasking() throws Exception {
         ctx.setDatabase("test");
-        new MockUp<StmtExecutor>() {
-            @Mock
-            public boolean isForwardToMaster() {
-                return false;
-            }
-        };
-        ctx.setEnv(env);
-        Config.enable_nereids_load = true;
-        // testing for https://github.com/apache/doris/issues/62140
-        String sql = "ALTER USER 'test_user62140'@'%' IDENTIFIED BY '123456'";
-        String res = "ALTER USER 'test_user62140'@'%' IDENTIFIED BY '*XXX'";
-        parseAndCheck(sql, res);
+        try (MockedConstruction<StmtExecutor> ignored = Mockito.mockConstruction(StmtExecutor.class,
+                Mockito.withSettings().defaultAnswer(Mockito.CALLS_REAL_METHODS),
+                (mock, context) -> {
+                    Mockito.doReturn(false).when(mock).isForwardToMaster();
+                    Profile profile = new Profile(false, 0, 0);
+                    Deencapsulation.setField(mock, "profile", profile);
+                    Mockito.doReturn(profile).when(mock).getProfile();
+                    Mockito.doReturn(ctx).when(mock).getContext();
+                    Mockito.doNothing().when(mock).execute();
+                    Deencapsulation.setField(mock, "context", ctx);
+                    if (context.arguments().size() >= 2
+                            && context.arguments().get(1) instanceof StatementBase) {
+                        Deencapsulation.setField(mock, "parsedStmt",
+                                context.arguments().get(1));
+                    }
+                    if (ctx.getStatementContext() == null) {
+                        ctx.setStatementContext(new StatementContext());
+                    }
+                })) {
+            ctx.setEnv(env);
+            ctx.setCurrentUserIdentity(UserIdentity.ROOT);
+            Config.enable_nereids_load = true;
+            // testing for https://github.com/apache/doris/issues/62140
+            String sql = "ALTER USER 'test_user62140'@'%' IDENTIFIED BY '123456'";
+            String res = "ALTER USER 'test_user62140'@'%' IDENTIFIED BY '*XXX'";
+            parseAndCheck(sql, res);
+        }
     }
 
     private void parseAndCheck(String sql, String expected) throws Exception {