Skip to content

Commit

Permalink
Fix unit test for coverage
Browse files Browse the repository at this point in the history
Signed-off-by: Tomoyuki Morita <[email protected]>
  • Loading branch information
ykmr1224 committed Nov 13, 2024
1 parent 0cfc9e2 commit 278a9c2
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 16 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -237,7 +237,8 @@ private void createIndex(String indexName) {
}
}

private long count(String indexName, QueryBuilder query) {
@VisibleForTesting
public long count(String indexName, QueryBuilder query) {
SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
searchSourceBuilder.query(query);
searchSourceBuilder.size(0);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@

package org.opensearch.sql.spark.leasemanager;

import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.Assert.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.when;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
Expand All @@ -23,22 +25,33 @@ class DefaultLeaseManagerTest {
@Mock private StateStore stateStore;

@Test
public void concurrentSessionRuleOnlyApplyToInteractiveQuery() {
assertTrue(
new DefaultLeaseManager.ConcurrentSessionRule(settings, stateStore)
.test(new LeaseRequest(JobType.BATCH, "mys3")));
assertTrue(
new DefaultLeaseManager.ConcurrentSessionRule(settings, stateStore)
.test(new LeaseRequest(JobType.STREAMING, "mys3")));
public void leaseManagerRejectsJobs() {
when(stateStore.count(any(), any())).thenReturn(3L);
when(settings.getSettingValue(any())).thenReturn(3);
DefaultLeaseManager defaultLeaseManager = new DefaultLeaseManager(settings, stateStore);

defaultLeaseManager.borrow(getLeaseRequest(JobType.BATCH));
assertThrows(ConcurrencyLimitExceededException.class, () ->
defaultLeaseManager.borrow(getLeaseRequest(JobType.INTERACTIVE)));
assertThrows(ConcurrencyLimitExceededException.class, () ->
defaultLeaseManager.borrow(getLeaseRequest(JobType.STREAMING)));
assertThrows(ConcurrencyLimitExceededException.class, () ->
defaultLeaseManager.borrow(getLeaseRequest(JobType.REFRESH)));
}

@Test
public void concurrentRefreshRuleNotAppliedToInteractiveAndBatchQuery() {
assertTrue(
new DefaultLeaseManager.ConcurrentRefreshJobRule(settings, stateStore)
.test(new LeaseRequest(JobType.INTERACTIVE, "mys3")));
assertTrue(
new DefaultLeaseManager.ConcurrentRefreshJobRule(settings, stateStore)
.test(new LeaseRequest(JobType.BATCH, "mys3")));
public void leaseManagerAcceptsJobs() {
when(stateStore.count(any(), any())).thenReturn(2L);
when(settings.getSettingValue(any())).thenReturn(3);
DefaultLeaseManager defaultLeaseManager = new DefaultLeaseManager(settings, stateStore);

defaultLeaseManager.borrow(getLeaseRequest(JobType.BATCH));
defaultLeaseManager.borrow(getLeaseRequest(JobType.INTERACTIVE));
defaultLeaseManager.borrow(getLeaseRequest(JobType.STREAMING));
defaultLeaseManager.borrow(getLeaseRequest(JobType.REFRESH));
}

private LeaseRequest getLeaseRequest(JobType jobType) {
return new LeaseRequest(jobType, "mys3");
}
}

0 comments on commit 278a9c2

Please sign in to comment.