Added helper function for SQLServiceTest

Signed-off-by: Guian Gumpac <guian.gumpac@improving.com>
Bit-Quill · Mar 10, 2023 · e0e2365 · e0e2365
1 parent 232196d
commit e0e2365
Showing 1 changed file with 30 additions and 58 deletions.
diff --git a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java
@@ -60,16 +60,9 @@ public void cleanup() throws InterruptedException {
     queryManager.awaitTermination(1, TimeUnit.SECONDS);
   }
 
-  @Test
-  public void canExecuteSqlQuery() {
-    doAnswer(invocation -> {
-      ResponseListener<QueryResponse> listener = invocation.getArgument(1);
-      listener.onResponse(new QueryResponse(schema, Collections.emptyList(), ""));
-      return null;
-    }).when(queryService).execute(any(), any());
-
+  private void execute(String query, String format) {
     sqlService.execute(
-        new SQLQueryRequest(new JSONObject(), "SELECT 123", QUERY, "jdbc"),
+        new SQLQueryRequest(new JSONObject(), query, QUERY, format),
         new ResponseListener<QueryResponse>() {
           @Override
           public void onResponse(QueryResponse response) {
@@ -83,16 +76,9 @@ public void onFailure(Exception e) {
         });
   }
 
-  @Test
-  public void canExecuteCsvFormatRequest() {
-    doAnswer(invocation -> {
-      ResponseListener<QueryResponse> listener = invocation.getArgument(1);
-      listener.onResponse(new QueryResponse(schema, Collections.emptyList(), ""));
-      return null;
-    }).when(queryService).execute(any(), any());
-
+  private void executeWithUnsupportedException(String query, String format) {
     sqlService.execute(
-        new SQLQueryRequest(new JSONObject(), "SELECT 123", QUERY, "csv"),
+        new SQLQueryRequest(new JSONObject(), query, QUERY, format),
         new ResponseListener<QueryResponse>() {
           @Override
           public void onResponse(QueryResponse response) {
@@ -101,66 +87,52 @@ public void onResponse(QueryResponse response) {
 
           @Override
           public void onFailure(Exception e) {
-            fail(e);
+            assert (e instanceof UnsupportedOperationException);
           }
         });
   }
 
   @Test
-  public void canExecuteJsonFormatRequest() {
+  public void canExecuteSqlQuery() {
     doAnswer(invocation -> {
       ResponseListener<QueryResponse> listener = invocation.getArgument(1);
       listener.onResponse(new QueryResponse(schema, Collections.emptyList(), ""));
       return null;
     }).when(queryService).execute(any(), any());
 
-    sqlService.execute(
-        new SQLQueryRequest(new JSONObject(), "SELECT FIELD FROM TABLE", QUERY, "json"),
-        new ResponseListener<QueryResponse>() {
-          @Override
-          public void onResponse(QueryResponse response) {
-            assertNotNull(response);
-          }
+    execute("SELECT 123", "jdbc");
+  }
 
-          @Override
-          public void onFailure(Exception e) {
-            fail(e);
-          }
-        });
+  @Test
+  public void canExecuteCsvFormatRequest() {
+    doAnswer(invocation -> {
+      ResponseListener<QueryResponse> listener = invocation.getArgument(1);
+      listener.onResponse(new QueryResponse(schema, Collections.emptyList(), ""));
+      return null;
+    }).when(queryService).execute(any(), any());
+
+    execute("SELECT 123", "csv");
   }
 
   @Test
-  public void canThrowUnsupportedExceptionWithUnsupportedJsonQuery() {
-    sqlService.execute(
-        new SQLQueryRequest(new JSONObject(), "SELECT CAST(FIELD AS DOUBLE)", QUERY, "json"),
-        new ResponseListener<QueryResponse>() {
-          @Override
-          public void onResponse(QueryResponse response) {
-            assertNotNull(response);
-          }
+  public void canExecuteJsonFormatRequest() {
+    doAnswer(invocation -> {
+      ResponseListener<QueryResponse> listener = invocation.getArgument(1);
+      listener.onResponse(new QueryResponse(schema, Collections.emptyList(), ""));
+      return null;
+    }).when(queryService).execute(any(), any());
 
-          @Override
-          public void onFailure(Exception e) {
-            assert (e instanceof UnsupportedOperationException);
-          }
-        });
+    execute("SELECT FIELD FROM TABLE", "json");
   }
 
   @Test
-  public void canThrowUnsupportedExceptionForHintsQuery() {
-    sqlService.execute(
-        new SQLQueryRequest(new JSONObject(), "SELECT /*! HINTS */ 123", QUERY, "json"),
-        new ResponseListener<QueryResponse>() {
-          @Override
-          public void onResponse(QueryResponse response) {
-            assertNotNull(response);
-          }
+  public void canThrowUnsupportedExceptionWithUnsupportedJsonQuery() {
+    executeWithUnsupportedException("SELECT CAST(FIELD AS DOUBLE)", "json");
+  }
 
-          @Override
-          public void onFailure(Exception e) {
-            assert (e instanceof UnsupportedOperationException);
-          }
-        });
+  @Test
+  public void canThrowUnsupportedExceptionForHintsQuery() {
+    executeWithUnsupportedException("SELECT /*! HINTS */ 123", "json");
   }
 
   @Test