package com.github.dockerjava.netty.handler;

import static com.github.dockerjava.netty.handler.HttpResponseStreamHandler.HttpResponseInputStream;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufInputStream;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;

import java.io.IOException;
import java.io.InputStream;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

import org.apache.commons.io.IOUtils;
import org.mockito.Mockito;
import org.testng.annotations.Test;

import com.github.dockerjava.core.async.ResultCallbackTemplate;

/**
 * @author Alexander Koshevoy
 */
public class HttpResponseStreamHandlerTest {
    @Test
    public void testNoBytesSkipped() throws Exception {
        ResultCallbackTest callback = new ResultCallbackTest();
        HttpResponseStreamHandler streamHandler = new HttpResponseStreamHandler(callback);
        ChannelHandlerContext ctx = Mockito.mock(ChannelHandlerContext.class);
        ByteBuf buffer = generateByteBuf();
        streamHandler.channelRead0(ctx, buffer);
        streamHandler.channelInactive(ctx);

        try (InputStream inputStream = callback.getInputStream()) {
            assertTrue(IOUtils.contentEquals(inputStream, new ByteBufInputStream(buffer)));
        }
    }

    @Test
    public void testReadByteByByte() throws Exception {
        ResultCallbackTest callback = new ResultCallbackTest();
        HttpResponseStreamHandler streamHandler = new HttpResponseStreamHandler(callback);
        ChannelHandlerContext ctx = Mockito.mock(ChannelHandlerContext.class);
        ByteBuf buffer = generateByteBuf();
        streamHandler.channelRead0(ctx, buffer);
        streamHandler.channelInactive(ctx);

        try (InputStream inputStream = callback.getInputStream()) {
            for (int i = 0; i < buffer.readableBytes(); i++) {
                int b = inputStream.read();
                assertEquals(b, buffer.getByte(i));
            }
            assertTrue(inputStream.read() == -1);
        }
    }

    @Test
    public void testCloseResponseStreamBeforeWrite() throws Exception {
        HttpResponseInputStream inputStream = new HttpResponseInputStream();
        ByteBuf buffer = generateByteBuf();

        inputStream.write(buffer);
        inputStream.close();
        inputStream.write(buffer);
    }

    @Test
    public void testCloseResponseStreamOnWrite() throws Exception {
        final HttpResponseInputStream inputStream = new HttpResponseInputStream();

        final ByteBuf buffer = generateByteBuf();

        final CountDownLatch firstWrite = new CountDownLatch(1);

        ExecutorService executor = Executors.newSingleThreadExecutor();
        Future<?> submit = executor.submit(new Runnable() {
            @Override
            public void run() {
                try {
                    inputStream.write(buffer);
                    firstWrite.countDown();
                    inputStream.write(buffer);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        });

        firstWrite.await();
        assertTrue(inputStream.available() > 0);

        // second write should have started
        Thread.sleep(500L);
        inputStream.close();

        submit.get();
    }

    @Test(expectedExceptions = IOException.class)
    public void testReadClosedResponseStream() throws Exception {
        HttpResponseInputStream inputStream = new HttpResponseInputStream();
        ByteBuf buffer = generateByteBuf();

        inputStream.write(buffer);
        inputStream.close();
        inputStream.read();
    }

    private ByteBuf generateByteBuf() {
        byte[] array = new byte[256];
        for (int i = 0; i < array.length; i++) {
            array[i] = (byte) i;
        }
        return Unpooled.copiedBuffer(array);
    }

    private static class ResultCallbackTest extends ResultCallbackTemplate<ResultCallbackTest, InputStream> {
        private InputStream stream;

        @Override
        public void onNext(InputStream stream) {
            this.stream = stream;
        }

        private InputStream getInputStream() {
            return stream;
        }
    }
}

@@ -604,24 +604,65 @@
 nt.Future;%0A%0A
+import io.netty.util.ReferenceCountUtil;%0A
 import org.a
@@ -1236,32 +1236,118 @@
 erateByteBuf();%0A
+        ByteBuf readBuffer = buffer.copy();%0A        assertEquals(buffer.refCnt(), 1);%0A
         streamHa
@@ -1355,33 +1355,32 @@
 dler.channelRead
-0
 (ctx, buffer);%0A 
@@ -1414,32 +1414,73 @@
 lInactive(ctx);%0A
+        assertEquals(buffer.refCnt(), 0);
 %0A        try (In
@@ -1613,17 +1613,21 @@
 tStream(
-b
+readB
 uffer)))
@@ -1630,32 +1630,80 @@
 r)));%0A        %7D%0A
+        ReferenceCountUtil.release(readBuffer);%0A
     %7D%0A%0A    @Test
@@ -2029,32 +2029,118 @@
 erateByteBuf();%0A
+        ByteBuf readBuffer = buffer.copy();%0A        assertEquals(buffer.refCnt(), 1);%0A
         streamHa
@@ -2156,17 +2156,16 @@
 nnelRead
-0
 (ctx, bu
@@ -2215,16 +2215,57 @@
 e(ctx);%0A
+        assertEquals(buffer.refCnt(), 0);
 %0A       
@@ -2357,17 +2357,21 @@
  0; i %3C 
-b
+readB
 uffer.re
@@ -2468,17 +2468,21 @@
 uals(b, 
-b
+readB
 uffer.ge
@@ -2562,24 +2562,72 @@
 ;%0A        %7D%0A
+        ReferenceCountUtil.release(readBuffer);%0A
     %7D%0A%0A    @

make sure ByteBuf is released