added assertion to check that records were read in while loop

ClickHouse · Aug 2, 2024 · f7235db · f7235db
1 parent 79a6589
commit f7235db
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java b/client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java
@@ -215,7 +215,7 @@ ClickHouseFormat[] getRowBinaryFormats() {
     @Test(groups = {"integration"}, dataProvider = "rowBinaryFormats")
     public void testRowBinaryQueries(ClickHouseFormat format)
             throws ExecutionException, InterruptedException {
-        final int rows = 1;
+        final int rows = 3;
         // TODO: replace with dataset with all primitive types of data
         // TODO: reusing same table name may lead to a conflict in tests?
 
@@ -229,11 +229,15 @@ public void testRowBinaryQueries(ClickHouseFormat format)
         ClickHouseBinaryFormatReader reader = createBinaryFormatReader(queryResponse, settings, tableSchema);
 
         Iterator<Map<String, Object>> dataIterator = data.iterator();
+        int rowsCount = 0;
         while (dataIterator.hasNext()) {
             Map<String, Object> expectedRecord = dataIterator.next();
             Map<String, Object> actualRecord = reader.next();
             Assert.assertEquals(actualRecord, expectedRecord);
+            rowsCount++;
         }
+
+        Assert.assertEquals(rowsCount, rows);
     }
 
     private static ClickHouseBinaryFormatReader createBinaryFormatReader(QueryResponse response, QuerySettings settings,
@@ -272,6 +276,7 @@ public void testBinaryStreamReader() throws Exception {
         schema.addColumn("col3", "String");
         schema.addColumn("host", "String");
         ClickHouseBinaryFormatReader reader = createBinaryFormatReader(queryResponse, settings, schema);
+        int rowsCount = 0;
         while (reader.next() != null) {
             String hostName = reader.readValue("host");
             Long col1 = reader.readValue("col1");
@@ -282,8 +287,9 @@ public void testBinaryStreamReader() throws Exception {
             Assert.assertEquals(reader.readValue(1), col1);
             Assert.assertEquals(reader.readValue(2), col3);
             Assert.assertEquals(reader.readValue(3), hostName);
+            rowsCount++;
         }
-
+        Assert.assertEquals(rowsCount, 10);
         Assert.assertFalse(reader.hasNext());
         Assert.assertNull(reader.next());
     }