Repository: netty/netty
File: codec-compression/src/test/java/io/netty/handler/codec/compression/ZstdEncoderTest.java (Java)

Code before:

        ByteBuf in = data.retainedSlice(written, 65535);
        assertTrue(channel.writeOutbound(in));

        ByteBuf in2 = data.retainedSlice(65535, dataLength - 65535);
        assertTrue(channel.writeOutbound(in2));

        assertTrue(channel.finish());

        ByteBuf decompressed = readDecompressed(dataLength);
        assertEquals(data, decompressed);

        decompressed.release();
        data.release();
    }

    @ParameterizedTest
    @MethodSource("smallData")
    public void testCompressionOfSmallBatchedFlow(final ByteBuf data) throws Exception {
        testCompressionOfBatchedFlow(data);
    }

    @ParameterizedTest
    @MethodSource("tinyData")
    public void testCompressionOfTinyData(final ByteBuf data) throws Exception {
        ZstdEncoder encoder = new ZstdEncoder();
        encoder.handlerAdded(ctx);
        VoidChannelPromise promise = new VoidChannelPromise(channel, true);
        encoder.write(ctx, data, promise);
        ArgumentCaptor<ByteBuf> outCaptor = ArgumentCaptor.forClass(ByteBuf.class);
        verify(ctx).write(outCaptor.capture(), eq(promise));

        ByteBuf out = outCaptor.getValue();
        assertThat(out.readableBytes()).isPositive();
        encoder.handlerRemoved(ctx);
        out.release();
    }

    @Override
    protected ByteBuf decompress(ByteBuf compressed, int originalLength) throws Exception {
        byte[] decompressed = new byte[originalLength];
        try (ZstdInputStream zstdIs = new ZstdInputStream(new ByteBufInputStream(compressed, true))) {

            int remaining = originalLength;
            while (remaining > 0) {
                int read = zstdIs.read(decompressed, originalLength - remaining, remaining);
                if (read > 0) {
                    remaining -= read;
                } else {
                    break;
                }

Review comment (suggestion):
Consider using `EmbeddedChannel` and not using mocks

Code after:
        ByteBuf decompressed = readDecompressed(dataLength);
        assertEquals(data, decompressed);

        decompressed.release();
        data.release();
    }

    @ParameterizedTest
    @MethodSource("smallData")
    public void testCompressionOfSmallBatchedFlow(final ByteBuf data) throws Exception {
        testCompressionOfBatchedFlow(data);
    }

    @Test
    public void testCompressionOfTinyData() throws Exception {
        ByteBuf data = Unpooled.copiedBuffer("Hello, World", CharsetUtil.UTF_8);
        assertTrue(channel.writeOutbound(data));
        assertTrue(channel.finish());

        ByteBuf out = channel.readOutbound();
        assertThat(out.readableBytes()).isPositive();
        out.release();
        assertNull(channel.readOutbound());
    }

    @Override
    protected ByteBuf decompress(ByteBuf compressed, int originalLength) throws Exception {
        byte[] decompressed = new byte[originalLength];
        try (ZstdInputStream zstdIs = new ZstdInputStream(new ByteBufInputStream(compressed, true))) {

            int remaining = originalLength;
            while (remaining > 0) {
                int read = zstdIs.read(decompressed, originalLength - remaining, remaining);
                if (read > 0) {
                    remaining -= read;
                } else {
                    break;
                }
            }
            assertEquals(-1, zstdIs.read());
        }

        return Unpooled.wrappedBuffer(decompressed);
    }
}