[KYUUBI #6160] Fix beeline test KyuubiBeeLineTest.testKyuubiBeelineCo…

…mment # 🔍 Description ## Issue References 🔗 Shell command `beeline` uses `KyuubiBeeline#initArgs`, and never uses `KyuubiBeeline#initArgsFromCliVars`. We should replace `KyuubiBeeline#initArgsFromCliVars` with `KyuubiBeeline#initArgs` in test `KyuubiBeeLineTest#testKyuubiBeelineComment`. ## Describe Your Solution 🔧 ## Types of changes 🔖 - [x] Bugfix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) ## Test Plan 🧪 #### Behavior Without This Pull Request ⚰️ #### Behavior With This Pull Request 🎉 #### Related Unit Tests --- # Checklist 📝 - [x] This patch was not authored or co-authored using [Generative Tooling](https://www.apache.org/legal/generative-tooling.html) **Be nice. Be informative.** Closes #6160 from zhouyifan279/beeline-comment-test. Closes #6160 c250af0 [zhouyifan279] Fix beeline test KyuubiBeeLineTest.testKyuubiBeelineComment 935aa37 [zhouyifan279] Fix flaky test KyuubiBeeLineTest.testKyuubiBeelineComment 29bf670 [zhouyifan279] Fix flaky test KyuubiBeeLineTest.testKyuubiBeelineComment f7a0742 [zhouyifan279] Fix flaky test KyuubiBeeLineTest.testKyuubiBeelineComment Authored-by: zhouyifan279 <[email protected]> Signed-off-by: Cheng Pan <[email protected]>
apache · Mar 11, 2024 · ca53942 · ca53942
1 parent 78e56af
commit ca53942
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 17 deletions.
diff --git a/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/BeeLine.java b/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/BeeLine.java
@@ -134,7 +134,7 @@ public class BeeLine implements Closeable {
   private final BeeLineOpts opts = new BeeLineOpts(this, System.getProperties());
   private String lastProgress = null;
   private final Map<SQLWarning, Date> seenWarnings = new HashMap<SQLWarning, Date>();
-  private final Commands commands = new Commands(this);
+  private Commands commands = new Commands(this);
   private OutputFile scriptOutputFile = null;
   private OutputFile recordOutputFile = null;
   private PrintStream outputStream = new PrintStream(System.out, true);
@@ -2382,4 +2382,8 @@ void setIsTestMode(boolean isTestMode) {
   boolean isTestMode() {
     return isTestMode;
   }
+
+  protected void setCommands(Commands commands) {
+    this.commands = commands;
+  }
 }
diff --git a/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/KyuubiBeeLine.java b/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/KyuubiBeeLine.java
@@ -48,7 +48,6 @@ public class KyuubiBeeLine extends BeeLine {
 
   public static final String KYUUBI_BEELINE_DEFAULT_JDBC_DRIVER =
       "org.apache.kyuubi.jdbc.KyuubiHiveDriver";
-  protected KyuubiCommands commands = new KyuubiCommands(this);
   private Driver defaultDriver;
 
   // copied from org.apache.hive.beeline.BeeLine
@@ -66,11 +65,7 @@ public KyuubiBeeLine() {
   @SuppressWarnings("deprecation")
   public KyuubiBeeLine(boolean isBeeLine) {
     super(isBeeLine);
-    try {
-      DynFields.builder().hiddenImpl(BeeLine.class, "commands").buildChecked(this).set(commands);
-    } catch (Throwable t) {
-      throw new ExceptionInInitializerError("Failed to inject kyuubi commands");
-    }
+    setCommands(new KyuubiCommands(this));
     try {
       defaultDriver =
           DynConstructors.builder()
@@ -299,4 +294,9 @@ int runInit() {
   boolean dispatch(String line) {
     return super.dispatch(isPythonMode() ? line : HiveStringUtils.removeComments(line));
   }
+
+  @Override
+  KyuubiCommands getCommands() {
+    return ((KyuubiCommands) super.getCommands());
+  }
 }
diff --git a/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/KyuubiDatabaseConnection.java b/kyuubi-hive-beeline/src/main/java/org/apache/hive/beeline/KyuubiDatabaseConnection.java
@@ -137,7 +137,7 @@ public Connection getConnectionFromDefaultDriver(String url, Properties properti
     eventNotifier.progressBarCompleted();
 
     Thread logThread =
-        new Thread(beeLine.commands.createLogRunnable(kyuubiConnection, eventNotifier));
+        new Thread(beeLine.getCommands().createLogRunnable(kyuubiConnection, eventNotifier));
     logThread.setDaemon(true);
     logThread.start();
     kyuubiConnection.setEngineLogThread(logThread);

diff --git a/kyuubi-hive-beeline/src/test/java/org/apache/hive/beeline/KyuubiBeeLineTest.java b/kyuubi-hive-beeline/src/test/java/org/apache/hive/beeline/KyuubiBeeLineTest.java
@@ -26,8 +26,13 @@
 import java.io.PrintStream;
 import org.apache.kyuubi.util.reflect.DynFields;
 import org.junit.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class KyuubiBeeLineTest {
+
+  private static final Logger LOG = LoggerFactory.getLogger(KyuubiBeeLineTest.class);
+
   @Test
   public void testKyuubiBeelineWithoutArgs() {
     KyuubiBeeLine kyuubiBeeLine = new KyuubiBeeLine();
@@ -90,15 +95,46 @@ public void testKyuubiBeeLinePythonMode() {
 
   @Test
   public void testKyuubiBeelineComment() {
-    KyuubiBeeLine kyuubiBeeLine = new KyuubiBeeLine();
-    int result = kyuubiBeeLine.initArgsFromCliVars(new String[] {"-e", "--comment show database;"});
-    assertEquals(0, result);
-    result = kyuubiBeeLine.initArgsFromCliVars(new String[] {"-e", "--comment\n show database;"});
-    assertEquals(1, result);
-    result =
-        kyuubiBeeLine.initArgsFromCliVars(
-            new String[] {"-e", "--comment line 1 \n    --comment line 2 \n show database;"});
-    assertEquals(1, result);
+    KyuubiBeeLine interceptedKyuubiBeeLine =
+        new KyuubiBeeLine() {
+          @Override
+          boolean dispatch(String line) {
+            if (line != null && line.startsWith("!connect")) {
+              LOG.info("Return true for command: {} to pretend connection is established.", line);
+              return true;
+            }
+            return super.dispatch(line);
+          }
+        };
+
+    String[] cmd = new String[] {""};
+    KyuubiCommands interceptedCommands =
+        new KyuubiCommands(interceptedKyuubiBeeLine) {
+          @Override
+          public boolean sql(String line, boolean entireLineAsCommand) {
+            LOG.info("Return true for sql {} to pretend sql is executed successfully.", line);
+            cmd[0] = line;
+            return true;
+          }
+        };
+    interceptedKyuubiBeeLine.setCommands(interceptedCommands);
+
+    interceptedKyuubiBeeLine.initArgs(
+        new String[] {"-u", "dummy_url", "-e", "--comment show database;"});
+    assertEquals(0, cmd[0].length());
+
+    // Beeline#exit must be false to execute sql
+    interceptedKyuubiBeeLine.setExit(false);
+    interceptedKyuubiBeeLine.initArgs(
+        new String[] {"-u", "dummy_url", "-e", "--comment\n show database;"});
+    assertEquals("show database;", cmd[0]);
+
+    interceptedKyuubiBeeLine.setExit(false);
+    interceptedKyuubiBeeLine.initArgs(
+        new String[] {
+          "-u", "dummy_url", "-e", "--comment line 1 \n    --comment line 2 \n show database;"
+        });
+    assertEquals("show database;", cmd[0]);
   }
 
   static class BufferPrintStream extends PrintStream {