Skip to content

Commit

Permalink
Add more tests
Browse files Browse the repository at this point in the history 
Signed-off-by: Louis Chu <clingzhi@amazon.com>
  • Loading branch information
@noCharger
noCharger committed Jul 18, 2024
1 parent bf0eba6 commit c61ccfa
Show file tree
Hide file tree
Showing 4 changed files with 102 additions and 26 deletions.
32 changes: 21 additions & 11 deletions async-query-core/src/main/antlr/SqlBaseParser.g4
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,11 +61,18 @@ compoundBody

compoundStatement
: statement
| setStatementWithOptionalVarKeyword
| beginEndCompoundBlock
;

setStatementWithOptionalVarKeyword
: SET (VARIABLE | VAR)? assignmentList #setVariableWithOptionalKeyword
| SET (VARIABLE | VAR)? LEFT_PAREN multipartIdentifierList RIGHT_PAREN EQ
LEFT_PAREN query RIGHT_PAREN #setVariableWithOptionalKeyword
;

singleStatement
: statement SEMICOLON* EOF
: (statement|setResetStatement) SEMICOLON* EOF
;

beginLabel
Expand Down Expand Up @@ -212,7 +219,7 @@ statement
identifierReference dataType? variableDefaultExpression? #createVariable
| DROP TEMPORARY VARIABLE (IF EXISTS)? identifierReference #dropVariable
| EXPLAIN (LOGICAL | FORMATTED | EXTENDED | CODEGEN | COST)?
statement #explain
(statement|setResetStatement) #explain
| SHOW TABLES ((FROM | IN) identifierReference)?
(LIKE? pattern=stringLit)? #showTables
| SHOW TABLE EXTENDED ((FROM | IN) ns=identifierReference)?
Expand Down Expand Up @@ -251,26 +258,29 @@ statement
| (MSCK)? REPAIR TABLE identifierReference
(option=(ADD|DROP|SYNC) PARTITIONS)? #repairTable
| op=(ADD | LIST) identifier .*? #manageResource
| SET COLLATION collationName=identifier #setCollation
| SET ROLE .*? #failNativeCommand
| CREATE INDEX (IF errorCapturingNot EXISTS)? identifier ON TABLE?
identifierReference (USING indexType=identifier)?
LEFT_PAREN columns=multipartIdentifierPropertyList RIGHT_PAREN
(OPTIONS options=propertyList)? #createIndex
| DROP INDEX (IF EXISTS)? identifier ON TABLE? identifierReference #dropIndex
| unsupportedHiveNativeCommands .*? #failNativeCommand
;

setResetStatement
: SET COLLATION collationName=identifier #setCollation
| SET ROLE .*? #failSetRole
| SET TIME ZONE interval #setTimeZone
| SET TIME ZONE timezone #setTimeZone
| SET TIME ZONE .*? #setTimeZone
| SET (VARIABLE | VAR) assignmentList #setVariable
| SET (VARIABLE | VAR) LEFT_PAREN multipartIdentifierList RIGHT_PAREN EQ
LEFT_PAREN query RIGHT_PAREN #setVariable
LEFT_PAREN query RIGHT_PAREN #setVariable
| SET configKey EQ configValue #setQuotedConfiguration
| SET configKey (EQ .*?)? #setConfiguration
| SET .*? EQ configValue #setQuotedConfiguration
| SET .*? #setConfiguration
| RESET configKey #resetQuotedConfiguration
| RESET .*? #resetConfiguration
| CREATE INDEX (IF errorCapturingNot EXISTS)? identifier ON TABLE?
identifierReference (USING indexType=identifier)?
LEFT_PAREN columns=multipartIdentifierPropertyList RIGHT_PAREN
(OPTIONS options=propertyList)? #createIndex
| DROP INDEX (IF EXISTS)? identifier ON TABLE? identifierReference #dropIndex
| unsupportedHiveNativeCommands .*? #failNativeCommand
;

executeImmediate
Expand Down
2 changes: 1 addition & 1 deletion ...query/src/main/java/org/opensearch/sql/spark/scheduler/OpenSearchAsyncQueryScheduler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ public void updateJob(OpenSearchRefreshIndexJobRequest request) throws IOExcepti
ActionFuture<UpdateResponse> updateResponseActionFuture = client.update(updateRequest);
updateResponse = updateResponseActionFuture.actionGet();
} catch (DocumentMissingException exception) {
throw new IllegalArgumentException("Job with name: " + request.getName() + " doesn't exist");
throw new IllegalArgumentException("Job: " + request.getName() + " doesn't exist");
} catch (Exception e) {
throw new RuntimeException(e);
}
Expand Down
62 changes: 52 additions & 10 deletions ...y/src/test/java/org/opensearch/sql/spark/scheduler/OpenSearchAsyncQuerySchedulerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import org.opensearch.client.Client;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.action.ActionFuture;
import org.opensearch.index.engine.DocumentMissingException;
import org.opensearch.index.engine.VersionConflictEngineException;
import org.opensearch.sql.spark.scheduler.model.OpenSearchRefreshIndexJobRequest;
import org.opensearch.threadpool.ThreadPool;
Expand All @@ -44,6 +45,8 @@ public class OpenSearchAsyncQuerySchedulerTest {

private static final String TEST_SCHEDULER_INDEX_NAME = "testQS";

private static final String TEST_JOB_ID = "testJobId";

@Mock(answer = Answers.RETURNS_DEEP_STUBS)
private Client client;

Expand Down Expand Up @@ -88,7 +91,7 @@ public void testScheduleJob() {

OpenSearchRefreshIndexJobRequest request =
OpenSearchRefreshIndexJobRequest.builder()
.jobName("testJob")
.jobName(TEST_JOB_ID)
.lastUpdateTime(Instant.now())
.build();

Expand All @@ -112,7 +115,7 @@ public void testScheduleJobWithExistingJob() {

OpenSearchRefreshIndexJobRequest request =
OpenSearchRefreshIndexJobRequest.builder()
.jobName("testJob")
.jobName(TEST_JOB_ID)
.lastUpdateTime(Instant.now())
.build();

Expand All @@ -131,30 +134,28 @@ public void testScheduleJobWithExistingJob() {

@Test
public void testUnscheduleJob() throws IOException {
String jobId = "testJob";

when(clusterService.state().routingTable().hasIndex(SCHEDULER_INDEX_NAME)).thenReturn(true);

when(updateResponseActionFuture.actionGet()).thenReturn(updateResponse);
when(updateResponse.getResult()).thenReturn(DocWriteResponse.Result.UPDATED);

when(client.update(any(UpdateRequest.class))).thenReturn(updateResponseActionFuture);

scheduler.unscheduleJob(jobId);
scheduler.unscheduleJob(TEST_JOB_ID);

ArgumentCaptor<UpdateRequest> captor = ArgumentCaptor.forClass(UpdateRequest.class);
verify(client).update(captor.capture());

UpdateRequest capturedRequest = captor.getValue();
assertEquals(jobId, capturedRequest.id());
assertEquals(TEST_JOB_ID, capturedRequest.id());
assertEquals(WriteRequest.RefreshPolicy.IMMEDIATE, capturedRequest.getRefreshPolicy());
}

@Test
public void testUpdateJob() throws IOException {
OpenSearchRefreshIndexJobRequest request =
OpenSearchRefreshIndexJobRequest.builder()
.jobName("testJob")
.jobName(TEST_JOB_ID)
.lastUpdateTime(Instant.now())
.build();

Expand All @@ -177,8 +178,6 @@ public void testUpdateJob() throws IOException {

@Test
public void testRemoveJob() {
String jobId = "testJob";

when(clusterService.state().routingTable().hasIndex(SCHEDULER_INDEX_NAME)).thenReturn(true);

DeleteResponse deleteResponse = mock(DeleteResponse.class);
Expand All @@ -187,7 +186,7 @@ public void testRemoveJob() {

when(client.delete(any(DeleteRequest.class))).thenReturn(deleteResponseActionFuture);

scheduler.removeJob(jobId);
scheduler.removeJob(TEST_JOB_ID);

ArgumentCaptor<DeleteRequest> captor = ArgumentCaptor.forClass(DeleteRequest.class);
verify(client).delete(captor.capture());
Expand Down Expand Up @@ -235,4 +234,47 @@ public void testCreateAsyncQuerySchedulerIndexFailure() throws IOException {
"Internal server error while creating .async-query-scheduler index: Error creating index",
exception.getMessage());
}

@Test
public void testUpdateJobNotFound() throws IOException {
OpenSearchRefreshIndexJobRequest request =
OpenSearchRefreshIndexJobRequest.builder()
.jobName(TEST_JOB_ID)
.lastUpdateTime(Instant.now())
.build();

when(clusterService.state().routingTable().hasIndex(SCHEDULER_INDEX_NAME)).thenReturn(true);

when(client.update(any(UpdateRequest.class)))
.thenThrow(new DocumentMissingException(null, null));

IllegalArgumentException exception =
assertThrows(
IllegalArgumentException.class,
() -> {
scheduler.updateJob(request);
});

assertEquals("Job: testJob doesn't exist", exception.getMessage());
}

@Test
public void testRemoveJobNotFound() {
when(clusterService.state().routingTable().hasIndex(SCHEDULER_INDEX_NAME)).thenReturn(true);

DeleteResponse deleteResponse = mock(DeleteResponse.class);
when(deleteResponseActionFuture.actionGet()).thenReturn(deleteResponse);
when(deleteResponse.getResult()).thenReturn(DocWriteResponse.Result.NOT_FOUND);

when(client.delete(any(DeleteRequest.class))).thenReturn(deleteResponseActionFuture);

IllegalArgumentException exception =
assertThrows(
IllegalArgumentException.class,
() -> {
scheduler.removeJob(TEST_JOB_ID);
});

assertEquals("Job : testJob doesn't exist", exception.getMessage());
}
}
32 changes: 28 additions & 4 deletions doctest/build.gradle
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,15 +145,39 @@ testClusters {
return new RegularFile() {
@Override
File getAsFile() {
File dir = new File('./doctest/' + jsPlugin)
// Use absolute paths
String basePath = new File('.').getCanonicalPath()
File dir = new File(basePath + File.separator + 'doctest' + File.separator + jsPlugin)

// Log the directory path for debugging
println("Creating directory: " + dir.getAbsolutePath())

// Create directory if it doesn't exist
if (!dir.exists()) {
dir.mkdirs()
if (!dir.mkdirs()) {
throw new IOException("Failed to create directory: " + dir.getAbsolutePath())
}
}

// Define the file path
File f = new File(dir, jsPlugin + '-' + opensearch_build + '.zip')

// Download file if it doesn't exist
if (!f.exists()) {
new URL(bwcOpenSearchJSDownload).withInputStream{ ins -> f.withOutputStream{ it << ins } }
println("Downloading file from: " + bwcOpenSearchJSDownload)
println("Saving to file: " + f.getAbsolutePath())

new URL(bwcOpenSearchJSDownload).withInputStream { ins ->
f.withOutputStream { it << ins }
}
}
return fileTree(jsPlugin).getSingleFile()

// Check if the file was created successfully
if (!f.exists()) {
throw new FileNotFoundException("File was not created: " + f.getAbsolutePath())
}

return fileTree(f.getParent()).matching { include f.getName() }.singleFile
}
}
}
Expand Down

0 comments on commit c61ccfa

Please sign in to comment.