Fix tests may fail in Proxy Backend (#11872)

* Fix tests may fail in Proxy Backend * Fix checkstyle * Add static import * Simplify code
apache · Aug 18, 2021 · 0ccdc93 · 0ccdc93
1 parent 6e6278f
commit 0ccdc93
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 4 deletions.
diff --git a/...nd/communication/jdbc/statement/impl/MySQLStatementMemoryStrictlyFetchSizeSetterTest.java b/...nd/communication/jdbc/statement/impl/MySQLStatementMemoryStrictlyFetchSizeSetterTest.java
@@ -17,28 +17,61 @@
 
 package org.apache.shardingsphere.proxy.backend.communication.jdbc.statement.impl;
 
+import lombok.SneakyThrows;
+import org.apache.shardingsphere.infra.config.properties.ConfigurationPropertyKey;
+import org.apache.shardingsphere.infra.context.manager.ContextManager;
+import org.apache.shardingsphere.proxy.backend.context.ProxyContext;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.junit.MockitoJUnitRunner;
 
+import java.lang.reflect.Field;
 import java.sql.SQLException;
 import java.sql.Statement;
 
 import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertThat;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+@RunWith(MockitoJUnitRunner.class)
 public final class MySQLStatementMemoryStrictlyFetchSizeSetterTest {
 
+    private static ContextManager originContextManager;
+
+    @BeforeClass
+    public static void setup() {
+        originContextManager = swapContextManager(mock(ContextManager.class, RETURNS_DEEP_STUBS));
+    }
+
     @Test
     public void assertSetFetchSize() throws SQLException {
+        when(ProxyContext.getInstance().getContextManager().getMetaDataContexts().getProps().<Integer>getValue(ConfigurationPropertyKey.PROXY_BACKEND_QUERY_FETCH_SIZE)).thenReturn(-1);
         Statement statement = mock(Statement.class);
         new MySQLStatementMemoryStrictlyFetchSizeSetter().setFetchSize(statement);
-        verify(statement, times(1)).setFetchSize(Integer.MIN_VALUE);
+        verify(statement).setFetchSize(Integer.MIN_VALUE);
     }
 
     @Test
     public void assertGetType() {
         assertThat(new MySQLStatementMemoryStrictlyFetchSizeSetter().getType(), is("MySQL"));
     }
+
+    @AfterClass
+    public static void tearDown() {
+        swapContextManager(originContextManager);
+    }
+
+    @SneakyThrows
+    private static ContextManager swapContextManager(final ContextManager newContextManager) {
+        Field contextManagerField = ProxyContext.class.getDeclaredField("contextManager");
+        contextManagerField.setAccessible(true);
+        ContextManager result = (ContextManager) contextManagerField.get(ProxyContext.getInstance());
+        contextManagerField.set(ProxyContext.getInstance(), newContextManager);
+        return result;
+    }
 }
diff --git a/...mmunication/jdbc/statement/impl/PostgreSQLStatementMemoryStrictlyFetchSizeSetterTest.java b/...mmunication/jdbc/statement/impl/PostgreSQLStatementMemoryStrictlyFetchSizeSetterTest.java
@@ -17,28 +17,61 @@
 
 package org.apache.shardingsphere.proxy.backend.communication.jdbc.statement.impl;
 
+import lombok.SneakyThrows;
+import org.apache.shardingsphere.infra.config.properties.ConfigurationPropertyKey;
+import org.apache.shardingsphere.infra.context.manager.ContextManager;
+import org.apache.shardingsphere.proxy.backend.context.ProxyContext;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.junit.MockitoJUnitRunner;
 
+import java.lang.reflect.Field;
 import java.sql.SQLException;
 import java.sql.Statement;
 
 import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertThat;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+@RunWith(MockitoJUnitRunner.class)
 public final class PostgreSQLStatementMemoryStrictlyFetchSizeSetterTest {
 
+    private static ContextManager originContextManager;
+
+    @BeforeClass
+    public static void setup() {
+        originContextManager = swapContextManager(mock(ContextManager.class, RETURNS_DEEP_STUBS));
+    }
+
     @Test
     public void assertSetFetchSize() throws SQLException {
+        when(ProxyContext.getInstance().getContextManager().getMetaDataContexts().getProps().<Integer>getValue(ConfigurationPropertyKey.PROXY_BACKEND_QUERY_FETCH_SIZE)).thenReturn(-1);
         Statement statement = mock(Statement.class);
         new PostgreSQLStatementMemoryStrictlyFetchSizeSetter().setFetchSize(statement);
-        verify(statement, times(1)).setFetchSize(1);
+        verify(statement).setFetchSize(1);
     }
 
     @Test
     public void assertGetType() {
         assertThat(new PostgreSQLStatementMemoryStrictlyFetchSizeSetter().getType(), is("PostgreSQL"));
     }
+
+    @AfterClass
+    public static void tearDown() {
+        swapContextManager(originContextManager);
+    }
+
+    @SneakyThrows
+    private static ContextManager swapContextManager(final ContextManager newContextManager) {
+        Field contextManagerField = ProxyContext.class.getDeclaredField("contextManager");
+        contextManagerField.setAccessible(true);
+        ContextManager result = (ContextManager) contextManagerField.get(ProxyContext.getInstance());
+        contextManagerField.set(ProxyContext.getInstance(), newContextManager);
+        return result;
+    }
 }