From 0e0292bff2bcadc73dc0c0ea547bc439c6270ba9 Mon Sep 17 00:00:00 2001
From: SHaaD94 <eugenzuy@gmail.com>
Date: Wed, 15 Jan 2025 11:25:40 +0700
Subject: [PATCH] Fix tests

Signed-off-by: Evgeniy Zuykin <evgeniy.zuykin@agoda.com>
---
 .../starrocks/server/SqlBlacklistTest.java    | 107 +++++++++---------
 1 file changed, 51 insertions(+), 56 deletions(-)

diff --git a/fe/fe-core/src/test/java/com/starrocks/server/SqlBlacklistTest.java b/fe/fe-core/src/test/java/com/starrocks/server/SqlBlacklistTest.java
index e648669040f31..5895509f083f9 100644
--- a/fe/fe-core/src/test/java/com/starrocks/server/SqlBlacklistTest.java
+++ b/fe/fe-core/src/test/java/com/starrocks/server/SqlBlacklistTest.java
@@ -15,6 +15,7 @@
 package com.starrocks.server;
 
 import com.starrocks.analysis.RedirectStatus;
+import com.starrocks.common.jmockit.Deencapsulation;
 import com.starrocks.meta.BlackListSql;
 import com.starrocks.meta.SqlBlackList;
 import com.starrocks.persist.DeleteSqlBlackLists;
@@ -32,11 +33,12 @@
 import mockit.Expectations;
 import mockit.Mock;
 import mockit.MockUp;
-import mockit.Mocked;
 import org.junit.Assert;
+import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.jupiter.api.Assertions;
+import org.mockito.Mockito;
 
 import java.util.List;
 import java.util.UUID;
@@ -45,40 +47,35 @@
 import static com.starrocks.sql.analyzer.AnalyzeTestUtil.parseSql;
 
 public class SqlBlacklistTest {
+    GlobalStateMgr state;
+    SqlBlackList sqlBlackList;
+    EditLog editLog;
+    ConnectContext connectContext;
+
     @BeforeClass
-    public static void beforeEach() throws Exception {
-        UtFrameUtils.createMinStarRocksCluster();
+    public static void beforeClass() throws Exception {
         AnalyzeTestUtil.init();
-        UtFrameUtils.setUpForPersistTest();
     }
 
-    @Test
-    public void testAddSQLBlacklist(@Mocked EditLog editLog) throws Exception {
-        GlobalStateMgr.getCurrentState().waitForReady();
-
-        SqlBlackList sqlBlackList = new SqlBlackList();
-
-        MockUp<GlobalStateMgr> mockUp = new MockUp<GlobalStateMgr>() {
-            @Mock
-            public SqlBlackList getSqlBlackList() {
-                return sqlBlackList;
-            }
-            @Mock
-            public EditLog getEditLog() {
-                return editLog;
-            }
-        };
+    @Before
+    public void beforeEach() {
+        state = Deencapsulation.newInstance(GlobalStateMgr.class);
+        sqlBlackList = new SqlBlackList();
+        connectContext = UtFrameUtils.createDefaultCtx();
+        editLog = Mockito.mock(EditLog.class);
+        connectContext.setQueryId(UUID.randomUUID());
+    }
 
-        new Expectations(editLog) {
+    @Test
+    public void testAddSQLBlacklist() throws Exception {
+        mockupGlobalState();
+        new Expectations() {
             {
                 editLog.logAddSQLBlackList(new SqlBlackListPersistInfo(0, ".+"));
                 times = 1;
             }
         };
 
-        ConnectContext connectContext = UtFrameUtils.createDefaultCtx();
-        connectContext.setQueryId(UUID.randomUUID());
-
         AddSqlBlackListStmt addStatement = (AddSqlBlackListStmt) parseSql("ADD SQLBLACKLIST \".+\";");
         Assert.assertEquals(addStatement.getSql(), ".+");
 
@@ -91,24 +88,11 @@ public EditLog getEditLog() {
     }
 
     @Test
-    public void testShowBlacklist() throws Exception {
-        GlobalStateMgr.getCurrentState().waitForReady();
-
-        SqlBlackList sqlBlackList = new SqlBlackList();
-
+    public void testShowBlacklist() {
+        mockupGlobalState();
         sqlBlackList.put(Pattern.compile("qwert"));
         sqlBlackList.put(Pattern.compile("abcde"));
 
-        MockUp<GlobalStateMgr> mockUp = new MockUp<GlobalStateMgr>() {
-            @Mock
-            public SqlBlackList getSqlBlackList() {
-                return sqlBlackList;
-            }
-        };
-
-        ConnectContext connectContext = UtFrameUtils.createDefaultCtx();
-        connectContext.setQueryId(UUID.randomUUID());
-
         ShowSqlBlackListStmt showSqlStatement = (ShowSqlBlackListStmt) parseSql("SHOW SQLBLACKLIST");
 
         ShowResultSet resultSet = ShowExecutor.execute(showSqlStatement, connectContext);
@@ -122,32 +106,18 @@ public SqlBlackList getSqlBlackList() {
     }
 
     @Test
-    public void testDeleteSqlBlacklist(@Mocked EditLog editLog) throws Exception {
-        SqlBlackList sqlBlackList = new SqlBlackList();
+    public void testDeleteSqlBlacklist() throws Exception {
+        mockupGlobalState();
         long id1 = sqlBlackList.put(Pattern.compile("qwert"));
         long id2 = sqlBlackList.put(Pattern.compile("abcde"));
 
-        MockUp<GlobalStateMgr> mockUp = new MockUp<GlobalStateMgr>() {
-            @Mock
-            public SqlBlackList getSqlBlackList() {
-                return sqlBlackList;
-            }
-            @Mock
-            public EditLog getEditLog() {
-                return editLog;
-            }
-        };
-
-        new Expectations(editLog) {
+        new Expectations() {
             {
                 editLog.logDeleteSQLBlackList(new DeleteSqlBlackLists(List.of(id1, id2)));
                 times = 1;
             }
         };
 
-        ConnectContext connectContext = UtFrameUtils.createDefaultCtx();
-        connectContext.setQueryId(UUID.randomUUID());
-
         StmtExecutor deleteStatementExecutor = new StmtExecutor(connectContext, new DelSqlBlackListStmt(List.of(id1, id2)));
         deleteStatementExecutor.execute();
         Assert.assertTrue(sqlBlackList
@@ -183,6 +153,8 @@ public void testSaveLoadBlackListImage() throws Exception {
 
     @Test
     public void testSqlBlacklistJournalOperations() throws Exception {
+        UtFrameUtils.createMinStarRocksCluster();
+        UtFrameUtils.setUpForPersistTest();
         UtFrameUtils.PseudoJournalReplayer.resetFollowerJournalQueue();
 
         // add blacklists
@@ -209,4 +181,27 @@ public void testSqlBlacklistJournalOperations() throws Exception {
 
     }
 
+    private void mockupGlobalState() {
+        MockUp<GlobalStateMgr> mockUp = new MockUp<GlobalStateMgr>() {
+            @Mock
+            GlobalStateMgr getCurrentState() {
+                return state;
+            }
+
+            @Mock
+            public SqlBlackList getSqlBlackList() {
+                return sqlBlackList;
+            }
+
+            @Mock
+            public boolean isLeader() {
+                return true;
+            }
+
+            @Mock
+            public EditLog getEditLog() {
+                return editLog;
+            }
+        };
+    }
 }