[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]>

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;
+  }
 }

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()

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 {