Address PR comments

Signed-off-by: Chen Dai <daichen@amazon.com>
dai-chen · Dec 1, 2023 · 52499af · 52499af
1 parent 0dc8b4d
commit 52499af
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 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
@@ -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
@@ -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);
     }
   }