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 index 683e443909..caafa96e8b 100644 --- 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;