Skip to content

Commit

Permalink
session module add UT (apache#10815)
Browse files Browse the repository at this point in the history
Co-authored-by: Cloudwise_Luke <282583553@qq.com>
  • Loading branch information
CloudWise-Lukemiao and 8669256 authored Aug 11, 2023
1 parent ff3239b commit 29d414b
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 44 deletions.
5 changes: 5 additions & 0 deletions iotdb-client/session/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,11 @@
<artifactId>mockito-core</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.powermock</groupId>
<artifactId>powermock-api-mockito2</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
<profiles>
<profile>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
package org.apache.iotdb.session.pool;

import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.isession.ISession;
import org.apache.iotdb.isession.SessionDataSet;
import org.apache.iotdb.isession.pool.ISessionPool;
import org.apache.iotdb.isession.pool.SessionDataSetWrapper;
Expand Down Expand Up @@ -48,9 +49,10 @@
import org.junit.FixMethodOrder;
import org.junit.Test;
import org.junit.runners.MethodSorters;
import org.mockito.ArgumentMatchers;
import org.mockito.Mock;
import org.mockito.Mockito;
import org.mockito.MockitoAnnotations;
import org.powermock.reflect.Whitebox;

import java.io.IOException;
import java.nio.ByteBuffer;
Expand All @@ -61,14 +63,12 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentLinkedDeque;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

@FixMethodOrder(MethodSorters.NAME_ASCENDING)
public class SessionPoolTest {
Expand Down Expand Up @@ -125,18 +125,25 @@ public void setUp() throws IoTDBConnectionException, StatementExecutionException
MockitoAnnotations.initMocks(this);
execResp.queryResult = FakedFirstFetchTsBlockResult();

when(client.executeStatementV2(any(TSExecuteStatementReq.class))).thenReturn(execResp);
when(execResp.getQueryId()).thenReturn(queryId);
when(execResp.getStatus()).thenReturn(successStatus);
Mockito.when(client.executeStatementV2(any(TSExecuteStatementReq.class))).thenReturn(execResp);
Mockito.when(execResp.getQueryId()).thenReturn(queryId);
Mockito.when(execResp.getStatus()).thenReturn(successStatus);

when(client.fetchMetadata(any(TSFetchMetadataReq.class))).thenReturn(fetchMetadataResp);
when(fetchMetadataResp.getStatus()).thenReturn(successStatus);
Mockito.when(client.fetchMetadata(any(TSFetchMetadataReq.class))).thenReturn(fetchMetadataResp);
Mockito.when(fetchMetadataResp.getStatus()).thenReturn(successStatus);

when(client.fetchResultsV2(any(TSFetchResultsReq.class))).thenReturn(fetchResultsResp);
when(fetchResultsResp.getStatus()).thenReturn(successStatus);
Mockito.when(client.fetchResultsV2(any(TSFetchResultsReq.class))).thenReturn(fetchResultsResp);
Mockito.when(fetchResultsResp.getStatus()).thenReturn(successStatus);

TSStatus closeResp = successStatus;
when(client.closeOperation(any(TSCloseOperationReq.class))).thenReturn(closeResp);
Mockito.when(client.closeOperation(any(TSCloseOperationReq.class))).thenReturn(closeResp);

sessionPool = new SessionPool("host", 11, "user", "password", 10);
ConcurrentLinkedDeque<ISession> queue = new ConcurrentLinkedDeque<>();
queue.add(session);

// 设置 SessionPool 对象的内部状态
Whitebox.setInternalState(sessionPool, "queue", queue);
}

@After
Expand All @@ -145,8 +152,21 @@ public void tearDown() {
sessionPool.close();
}

@Test(expected = StatementExecutionException.class)
public void testInsertAlignedTabletsWithStatementExecutionException()
throws IoTDBConnectionException, StatementExecutionException {
ISessionPool sessionPool = Mockito.mock(ISessionPool.class);
Map<String, Tablet> tablets = new HashMap<>();
boolean sorted = true;
doThrow(new StatementExecutionException(""))
.when(sessionPool)
.insertAlignedTablets(tablets, sorted);
sessionPool.insertAlignedTablets(tablets, sorted);
}

@Test
public void testInsertRecords() throws IoTDBConnectionException, StatementExecutionException {
public void testInsertRecords() throws Exception {
// 调用 insertRecords 方法
List<String> deviceIds = Arrays.asList("device1", "device2");
List<Long> timeList = Arrays.asList(1L, 2L);
List<List<String>> measurementsList =
Expand All @@ -158,33 +178,15 @@ public void testInsertRecords() throws IoTDBConnectionException, StatementExecut
Arrays.asList(TSDataType.DOUBLE, TSDataType.DOUBLE));
List<List<Object>> valuesList =
Arrays.asList(Arrays.asList(25.0f, 50.0f), Arrays.asList(220.0, 1.5));

sessionPool.insertRecords(deviceIds, timeList, measurementsList, typesList, valuesList);

verify(sessionPool, times(1))
.insertRecords(
ArgumentMatchers.eq(deviceIds),
ArgumentMatchers.eq(timeList),
ArgumentMatchers.eq(measurementsList),
ArgumentMatchers.eq(typesList),
ArgumentMatchers.eq(valuesList));
}

@Test(expected = StatementExecutionException.class)
public void testInsertAlignedTabletsWithStatementExecutionException()
throws IoTDBConnectionException, StatementExecutionException {
Map<String, Tablet> tablets = new HashMap<>();
boolean sorted = true;
doThrow(new StatementExecutionException(""))
.when(sessionPool)
.insertAlignedTablets(tablets, sorted);
sessionPool.insertAlignedTablets(tablets, sorted);
assertEquals(
1,
((ConcurrentLinkedDeque<ISession>) Whitebox.getInternalState(sessionPool, "queue")).size());
}

@Test
public void testExecuteQueryStatement()
throws IoTDBConnectionException, StatementExecutionException {
// 构造测试数据
String testSql = "select s2,s1,s0,s2 from root.vehicle.d0 ";

List<String> columns = new ArrayList<>();
Expand All @@ -199,14 +201,14 @@ public void testExecuteQueryStatement()
dataTypeList.add("INT32");
dataTypeList.add("FLOAT");

when(execResp.isSetColumns()).thenReturn(true);
when(execResp.getColumns()).thenReturn(columns);
when(execResp.isSetDataTypeList()).thenReturn(true);
when(execResp.getDataTypeList()).thenReturn(dataTypeList);
when(execResp.isSetOperationType()).thenReturn(true);
when(execResp.getOperationType()).thenReturn("QUERY");
when(execResp.isSetQueryId()).thenReturn(true);
when(execResp.getQueryId()).thenReturn(queryId);
Mockito.when(execResp.isSetColumns()).thenReturn(true);
Mockito.when(execResp.getColumns()).thenReturn(columns);
Mockito.when(execResp.isSetDataTypeList()).thenReturn(true);
Mockito.when(execResp.getDataTypeList()).thenReturn(dataTypeList);
Mockito.when(execResp.isSetOperationType()).thenReturn(true);
Mockito.when(execResp.getOperationType()).thenReturn("QUERY");
Mockito.when(execResp.isSetQueryId()).thenReturn(true);
Mockito.when(execResp.getQueryId()).thenReturn(queryId);

SessionDataSet sessionDataSet =
new SessionDataSet(
Expand All @@ -223,8 +225,8 @@ public void testExecuteQueryStatement()
10,
true,
10);
when(sessionPool.executeQueryStatement(any(String.class)))
.thenReturn(new SessionDataSetWrapper(sessionDataSet, session, sessionPool));

Mockito.when(session.executeQueryStatement(any(String.class))).thenReturn(sessionDataSet);

SessionDataSetWrapper sessionDataSetWrapper = sessionPool.executeQueryStatement(testSql);
List<String> columnNames = sessionDataSetWrapper.getColumnNames();
Expand Down

0 comments on commit 29d414b

Please sign in to comment.