diff --git a/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp b/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp
index 7a04de0324948..45776cb5f1a61 100644
--- a/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp
+++ b/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp
@@ -415,7 +415,7 @@ void AggregationTestBase::testAggregationsWithCompanion(
         toPlanStats(task->taskStats()).at(partialNodeId).inputRows;
     const auto finalInputRows =
         toPlanStats(task->taskStats()).at(finalNodeId).inputRows;
-    if (partialInputRows > 1) {
+    if (exec::injectedSpillCount() > 0) {
       EXPECT_LT(0, spilledBytes(*task))
           << "partial inputRows: " << partialInputRows
           << " final inputRows: " << finalInputRows
@@ -854,7 +854,7 @@ void AggregationTestBase::testAggregationsImpl(
         toPlanStats(task->taskStats()).at(partialNodeId).inputRows;
     const auto finalInputRows =
         toPlanStats(task->taskStats()).at(finalNodeId).inputRows;
-    if (partialInputRows > 1) {
+    if (exec::injectedSpillCount() > 0) {
       EXPECT_LT(0, spilledBytes(*task))
           << "partial inputRows: " << partialInputRows
           << " final inputRows: " << finalInputRows