Refactor test code

Signed-off-by: Chen Dai <[email protected]>
opensearch-project · Oct 5, 2023 · ff192ed · ff192ed
1 parent 4e78536
commit ff192ed
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 13 deletions.
diff --git a/.../src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImplTest.java b/.../src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImplTest.java
@@ -6,6 +6,7 @@
 package org.opensearch.sql.spark.asyncquery;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
@@ -112,14 +113,11 @@ void testCreateAsyncQueryWithExtraSparkSubmitParameter() {
 
     verify(sparkQueryDispatcher, times(1))
         .dispatch(
-            new DispatchQueryRequest(
-                "00fd775baqpu4g0p",
-                "select * from my_glue.default.http_logs",
-                "my_glue",
-                LangType.SQL,
-                "arn:aws:iam::270824043731:role/emr-job-execution-role",
-                TEST_CLUSTER_NAME,
-                "--conf spark.dynamicAllocation.enabled=false"));
+            argThat(
+                actualReq ->
+                    actualReq
+                        .getExtraSparkSubmitParams()
+                        .equals("--conf spark.dynamicAllocation.enabled=false")));
   }
 
   @Test

diff --git a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/model/SparkSubmitParametersTest.java b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/model/SparkSubmitParametersTest.java
@@ -25,6 +25,6 @@ public void testBuildWithExtraParameters() {
         SparkSubmitParameters.Builder.builder().extraParameters("--conf A=1").build().toString();
 
     // Assert the conf is included with a space
-    assertTrue(params.contains(" --conf A=1"));
+    assertTrue(params.endsWith(" --conf A=1"));
   }
 }
diff --git a/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java b/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java
@@ -698,21 +698,23 @@ void testDispatchQueryWithExtraSparkSubmitParameters() {
 
     String extraParameters = "--conf spark.dynamicAllocation.enabled=false";
     DispatchQueryRequest[] requests = {
-      constructDispatchQueryRequest( // SQL direct query
+      // SQL direct query
+      constructDispatchQueryRequest(
           "select * from my_glue.default.http_logs", LangType.SQL, extraParameters),
-      constructDispatchQueryRequest( // SQL index query
+      // SQL index query
+      constructDispatchQueryRequest(
           "create skipping index on my_glue.default.http_logs (status VALUE_SET)",
           LangType.SQL,
           extraParameters),
-      constructDispatchQueryRequest( // PPL query
+      // PPL query
+      constructDispatchQueryRequest(
           "source = my_glue.default.http_logs", LangType.PPL, extraParameters)
     };
 
     for (DispatchQueryRequest request : requests) {
       when(emrServerlessClient.startJobRun(any())).thenReturn(EMR_JOB_ID);
       sparkQueryDispatcher.dispatch(request);
 
-      // This test is only interested in Spark submit parameters
       verify(emrServerlessClient, times(1))
           .startJobRun(
               argThat(