Address PR comments

Signed-off-by: Chen Dai <daichen@amazon.com>
dai-chen · Dec 1, 2023 · 0dc8b4d · 0dc8b4d
1 parent 3582e4f
commit 0dc8b4d
Showing 1 changed file with 7 additions and 12 deletions.
diff --git a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryGetResultSpecTest.java b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryGetResultSpecTest.java
@@ -51,7 +51,7 @@ public void doSetUp() {
   @Test
   public void testInteractiveQueryGetResult() {
     createAsyncQuery("SELECT 1")
-        .withoutInteraction()
+        .withInteraction(InteractionStep::pluginSearchQueryResult)
         .assertQueryResults("waiting", null)
         .withInteraction(
             interaction -> {
@@ -65,7 +65,7 @@ public void testInteractiveQueryGetResult() {
   @Test
   public void testInteractiveQueryGetResultWithConcurrentEmrJobUpdate() {
     createAsyncQuery("SELECT 1")
-        .withoutInteraction()
+        .withInteraction(InteractionStep::pluginSearchQueryResult)
         .assertQueryResults("waiting", null)
         .withInteraction(
             interaction -> {
@@ -75,7 +75,7 @@ public void testInteractiveQueryGetResultWithConcurrentEmrJobUpdate() {
               return result;
             })
         .assertQueryResults("running", null)
-        .withoutInteraction()
+        .withInteraction(InteractionStep::pluginSearchQueryResult)
         .assertQueryResults("SUCCESS", ImmutableList.of(tupleValue(Map.of("1", 1))));
   }
 
@@ -102,7 +102,7 @@ public void testBatchQueryGetResultWithConcurrentEmrJobUpdate() {
               return result;
             })
         .assertQueryResults("running", null)
-        .withoutInteraction()
+        .withInteraction(InteractionStep::pluginSearchQueryResult)
         .assertQueryResults("SUCCESS", ImmutableList.of());
   }
 
@@ -145,7 +145,7 @@ public void testStreamingQueryGetResultWithConcurrentEmrJobUpdate() {
                 return result;
               })
           .assertQueryResults("running", null)
-          .withoutInteraction()
+          .withInteraction(InteractionStep::pluginSearchQueryResult)
           .assertQueryResults("SUCCESS", ImmutableList.of());
     } finally {
       mockIndex.deleteIndex();
@@ -162,7 +162,7 @@ public void testDropIndexQueryGetResult() {
     LocalEMRSClient emrClient = new LocalEMRSClient();
     emrClient.setJobState("Cancelled");
     createAsyncQuery(mockIndex.query, emrClient)
-        .withoutInteraction()
+        .withInteraction(InteractionStep::pluginSearchQueryResult)
         .assertQueryResults("SUCCESS", ImmutableList.of());
   }
 
@@ -177,7 +177,7 @@ public void testDropIndexQueryGetResultWithResultDocRefreshDelay() {
     createAsyncQuery(mockIndex.query, emrClient)
         .withInteraction(interaction -> new JSONObject()) // simulate result index refresh delay
         .assertQueryResults("running", null)
-        .withoutInteraction()
+        .withInteraction(InteractionStep::pluginSearchQueryResult)
         .assertQueryResults("SUCCESS", ImmutableList.of());
   }
 
@@ -219,11 +219,6 @@ public JSONObject getResultWithQueryId(String queryId, String resultIndex) {
               new CreateAsyncQueryRequest(query, DATASOURCE, LangType.SQL, null));
     }
 
-    AssertionHelper withoutInteraction() {
-      // No interaction with EMR-S job. Plugin searches query result only.
-      return withInteraction(InteractionStep::pluginSearchQueryResult);
-    }
-
     AssertionHelper withInteraction(Interaction interaction) {
       this.interaction = interaction;
       return this;