Improve GroupByHashYieldAssertion

trinodb · May 27, 2022 · 1cf9bc6 · 1cf9bc6
1 parent aa26c96
commit 1cf9bc6
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/core/trino-main/src/test/java/io/trino/operator/GroupByHashYieldAssertion.java b/core/trino-main/src/test/java/io/trino/operator/GroupByHashYieldAssertion.java
@@ -127,7 +127,10 @@ public static GroupByHashYieldResult finishOperatorWithYieldingGroupByHash(List<
                 // free the pool for the next iteration
                 memoryPool.free(anotherTaskId, "test", reservedMemoryInBytes);
                 // this required in case input is blocked
-                operator.getOutput();
+                output = operator.getOutput();
+                if (output != null) {
+                    result.add(output);
+                }
                 continue;
             }
 
@@ -140,7 +143,7 @@ public static GroupByHashYieldResult finishOperatorWithYieldingGroupByHash(List<
                 assertTrue(operator.getOperatorContext().isWaitingForMemory().isDone());
 
                 // assert the hash capacity is not changed; otherwise, we should have yielded
-                assertTrue(oldCapacity == getHashCapacity.apply(operator));
+                assertEquals((int) getHashCapacity.apply(operator), oldCapacity);
 
                 // We are not going to rehash; therefore, assert the memory increase only comes from the aggregator
                 assertLessThan(actualIncreasedMemory, additionalMemoryInBytes);