From 52499afeb7de50d3f85e4a589fdab4d823f49381 Mon Sep 17 00:00:00 2001 From: Chen Dai Date: Fri, 1 Dec 2023 11:26:37 -0800 Subject: [PATCH] Address PR comments Signed-off-by: Chen Dai --- .../spark/asyncquery/AsyncQueryExecutorServiceSpec.java | 7 ++++--- .../sql/spark/asyncquery/AsyncQueryGetResultSpecTest.java | 6 +++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java index ef96a0d6da..c7054dd200 100644 --- a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java +++ b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java @@ -15,6 +15,7 @@ import com.amazonaws.services.emrserverless.model.CancelJobRunResult; import com.amazonaws.services.emrserverless.model.GetJobRunResult; import com.amazonaws.services.emrserverless.model.JobRun; +import com.amazonaws.services.emrserverless.model.JobRunState; import com.google.common.base.Charsets; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -228,7 +229,7 @@ public static class LocalEMRSClient implements EMRServerlessClient { private int startJobRunCalled = 0; private int cancelJobRunCalled = 0; private int getJobResult = 0; - private String jobState = "RUNNING"; + private JobRunState jobState = JobRunState.RUNNING; @Getter private StartJobRequest jobRequest; @@ -243,7 +244,7 @@ public String startJobRun(StartJobRequest startJobRequest) { public GetJobRunResult getJobRunResult(String applicationId, String jobId) { getJobResult++; JobRun jobRun = new JobRun(); - jobRun.setState(jobState); + jobRun.setState(jobState.toString()); return new GetJobRunResult().withJobRun(jobRun); } @@ -265,7 +266,7 @@ public void getJobRunResultCalled(int expectedTimes) { assertEquals(expectedTimes, getJobResult); } - public void setJobState(String jobState) { + public void setJobState(JobRunState jobState) { this.jobState = jobState; } } 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 caafa96e8b..bba38693cd 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 @@ -160,7 +160,7 @@ public void testDropIndexQueryGetResult() { mockIndexState.refreshing(); LocalEMRSClient emrClient = new LocalEMRSClient(); - emrClient.setJobState("Cancelled"); + emrClient.setJobState(JobRunState.CANCELLED); createAsyncQuery(mockIndex.query, emrClient) .withInteraction(InteractionStep::pluginSearchQueryResult) .assertQueryResults("SUCCESS", ImmutableList.of()); @@ -173,7 +173,7 @@ public void testDropIndexQueryGetResultWithResultDocRefreshDelay() { mockIndexState.refreshing(); LocalEMRSClient emrClient = new LocalEMRSClient(); - emrClient.setJobState("Cancelled"); + emrClient.setJobState(JobRunState.CANCELLED); createAsyncQuery(mockIndex.query, emrClient) .withInteraction(interaction -> new JSONObject()) // simulate result index refresh delay .assertQueryResults("running", null) @@ -277,7 +277,7 @@ void emrJobUpdateStatementState(StatementState newState) { } void emrJobUpdateJobState(JobRunState jobState) { - emrClient.setJobState(jobState.toString()); + emrClient.setJobState(jobState); } }