agavra commented on code in PR #9767: URL: https://github.com/apache/pinot/pull/9767#discussion_r1018255056
########## pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java: ########## @@ -18,26 +18,208 @@ */ package org.apache.pinot.query.runtime.operator; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import java.util.Arrays; import java.util.List; import org.apache.calcite.sql.SqlKind; +import org.apache.pinot.common.datablock.DataBlock; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.common.Operator; import org.apache.pinot.core.operator.BaseOperator; import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.runtime.blocks.TransferableBlock; +import org.apache.pinot.query.runtime.blocks.TransferableBlockUtils; import org.apache.pinot.spi.data.FieldSpec; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; import org.testng.Assert; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.DOUBLE; +import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.INT; public class AggregateOperatorTest { + private AutoCloseable _mocks; + + @Mock + private Operator<TransferableBlock> _input; + + @BeforeMethod + public void setUp() { + _mocks = MockitoAnnotations.openMocks(this); + } + + @AfterMethod + public void tearDown() + throws Exception { + _mocks.close(); + } + + @Test + public void shouldHandleUpstreamErrorBlocks() { + // Given: + List<RexExpression> calls = ImmutableList.of(getSum(new RexExpression.InputRef(1))); + List<RexExpression> group = ImmutableList.of(new RexExpression.InputRef(0)); + + Mockito.when(_input.nextBlock()) + .thenReturn(TransferableBlockUtils.getErrorTransferableBlock(new Exception("foo!"))); + + DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); + AggregateOperator operator = new AggregateOperator(_input, outSchema, calls, group); + + // When: + TransferableBlock block1 = operator.nextBlock(); // build + + // Then: + Mockito.verify(_input, Mockito.times(1)).nextBlock(); + Assert.assertTrue(block1.isErrorBlock(), "Input errors should propagate immediately"); + } + + @Test + public void shouldHandleEndOfStreamBlockWithNoOtherInputs() { + // Given: + List<RexExpression> calls = ImmutableList.of(getSum(new RexExpression.InputRef(1))); + List<RexExpression> group = ImmutableList.of(new RexExpression.InputRef(0)); + + Mockito.when(_input.nextBlock()) + .thenReturn(TransferableBlockUtils.getEndOfStreamTransferableBlock()); + + DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); + AggregateOperator operator = new AggregateOperator(_input, outSchema, calls, group); + + // When: + TransferableBlock block = operator.nextBlock(); + + // Then: + Mockito.verify(_input, Mockito.times(1)).nextBlock(); + Assert.assertTrue(block.isEndOfStreamBlock(), "EOS blocks should propagate"); + } + + @Test + public void shouldHandleUpstreamNoOpBlocksWhileConstructing() { + // Given: + List<RexExpression> calls = ImmutableList.of(getSum(new RexExpression.InputRef(1))); + List<RexExpression> group = ImmutableList.of(new RexExpression.InputRef(0)); + + DataSchema inSchema = new DataSchema(new String[]{"group", "arg"}, new ColumnDataType[]{INT, INT}); + Mockito.when(_input.nextBlock()) + .thenReturn(block(inSchema, new Object[]{1, 1})) + .thenReturn(TransferableBlockUtils.getNoOpTransferableBlock()); + + DataSchema outSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); + AggregateOperator operator = new AggregateOperator(_input, outSchema, calls, group); Review Comment: > but i am also ok if the goal is to keep these tests isolated and easy to maintain. that's my goal :) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org