v2 changes

Trianz-Akshay · Sep 6, 2024 · 9d31d27 · 9d31d27
1 parent e866ed6
commit 9d31d27
Showing 1 changed file with 27 additions and 11 deletions.
diff --git a/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/TestUtils.java b/athena-timestream/src/test/java/com/amazonaws/athena/connectors/timestream/TestUtils.java
@@ -55,18 +55,34 @@ public static QueryResponse makeMockQueryResult(Schema schemaForRead, int numRow
 
     public static QueryResponse makeMockQueryResult(Schema schemaForRead, int numRows, int maxDataGenerationRow, boolean isRandomAZ)
     {
+        QueryResponse mockResult = mock(QueryResponse.class);
         final AtomicLong nextToken = new AtomicLong(0);
-        List<Row> rows = new ArrayList<>();
-        for (int i = 0; i < maxDataGenerationRow; i++) {
-            nextToken.incrementAndGet();
-            List<Datum> columnData = new ArrayList<>();
-            for (Field nextField : schemaForRead.getFields()) {
-                columnData.add(makeValue(nextField, i, isRandomAZ));
-            }
-            Row row = Row.builder().data(columnData).build();
-            rows.add(row);
-        }
-        return QueryResponse.builder().rows(rows).nextToken(nextToken.get() < numRows ? String.valueOf(nextToken) : null).build();
+
+        when(mockResult.rows()).thenAnswer((Answer<List<Row>>) invocationOnMock -> {
+                    List<Row> rows = new ArrayList<>();
+                    for (int i = 0; i < maxDataGenerationRow; i++) {
+                        nextToken.incrementAndGet();
+                        List<Datum> columnData = new ArrayList<>();
+                        for (Field nextField : schemaForRead.getFields()) {
+                            columnData.add(makeValue(nextField, i, isRandomAZ));
+                        }
+
+                        Row row = Row.builder().data(columnData).build();
+                        rows.add(row);
+                    }
+                    return rows;
+                }
+        );
+
+        when(mockResult.nextToken()).thenAnswer((Answer<String>) invocationOnMock -> {
+                    if (nextToken.get() < numRows) {
+                        return String.valueOf(nextToken.get());
+                    }
+                    return null;
+                }
+        );
+
+        return mockResult;
     }
 
     public static Datum makeValue(Field field, int num, boolean isRandomAZ)