package com.github.bedrin.jdbc.sniffer.servlet;

import com.github.bedrin.jdbc.sniffer.BaseTest;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;
import org.springframework.mock.web.MockHttpServletResponse;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import java.io.IOException;

import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class SnifferFilterTest extends BaseTest {

    @Mock
    private HttpServletRequest httpServletRequest;
    @Mock
    private HttpServletResponse httpServletResponse;
    @Mock
    private FilterChain filterChain;

    @Test
    public void testFilterNoQueries() throws IOException, ServletException {

        SnifferFilter filter = new SnifferFilter();

        filter.doFilter(httpServletRequest, httpServletResponse, filterChain);

        verify(httpServletResponse).addIntHeader(SnifferFilter.HEADER_NAME, 0);

    }

    @Test
    public void testFilterOneQuery() throws IOException, ServletException {

        doAnswer(invocation -> {executeStatement(); return null;}).
                when(filterChain).doFilter(any(), any());

        SnifferFilter filter = new SnifferFilter();

        filter.doFilter(httpServletRequest, httpServletResponse, filterChain);

        verify(httpServletResponse).addIntHeader(SnifferFilter.HEADER_NAME, 1);

    }

    @Test
    public void testFilterOneQueryWithOutput() throws IOException, ServletException {

        MockHttpServletResponse httpServletResponse = new MockHttpServletResponse();

        doAnswer(invocation -> {
            HttpServletResponse response = (HttpServletResponse) invocation.getArguments()[1];
            response.getOutputStream().write("Hello, World".getBytes());
            executeStatement();
            return null;
        }).when(filterChain).doFilter(any(), any());

        SnifferFilter filter = new SnifferFilter();

        filter.doFilter(httpServletRequest, httpServletResponse, filterChain);

        assertEquals(1, httpServletResponse.getHeaderValue(SnifferFilter.HEADER_NAME));
        assertArrayEquals("Hello, World".getBytes(), httpServletResponse.getContentAsByteArray());

    }

}
@@ -218,16 +218,76 @@
 Runner;%0A
+import org.springframework.mock.web.MockHttpServletRequest;%0A
 import o
@@ -331,32 +331,32 @@
 ervletResponse;%0A
-
 %0Aimport javax.se
@@ -417,54 +417,8 @@
 on;%0A
-import javax.servlet.http.HttpServletRequest;%0A
 impo
@@ -753,132 +753,8 @@
  %7B%0A%0A
-    @Mock%0A    private HttpServletRequest httpServletRequest;%0A    @Mock%0A    private HttpServletResponse httpServletResponse;%0A
     
@@ -877,32 +877,200 @@
 letException %7B%0A%0A
+        MockHttpServletResponse httpServletResponse = new MockHttpServletResponse();%0A        MockHttpServletRequest httpServletRequest = new MockHttpServletRequest();%0A%0A
         SnifferF
@@ -1186,39 +1186,48 @@
 hain);%0A%0A        
-verify(
+assertEquals(0, 
 httpServletRespo
@@ -1225,38 +1225,39 @@
 vletResponse
-).addIn
+.ge
 tHeader
+Value
 (SnifferFilt
@@ -1274,11 +1274,9 @@
 NAME
-, 0
+)
 );%0A%0A
@@ -1361,32 +1361,200 @@
 letException %7B%0A%0A
+        MockHttpServletResponse httpServletResponse = new MockHttpServletResponse();%0A        MockHttpServletRequest httpServletRequest = new MockHttpServletRequest();%0A%0A
         doAnswer
@@ -1809,15 +1809,24 @@
     
-verify(
+assertEquals(1, 
 http
@@ -1844,22 +1844,23 @@
 onse
-).addIn
+.ge
 tHeader
+Value
 (Sni
@@ -1885,11 +1885,9 @@
 NAME
-, 1
+)
 );%0A%0A
@@ -1982,32 +1982,32 @@
 letException %7B%0A%0A
-
         MockHttp
@@ -2074,16 +2074,98 @@
 ponse();
+%0A        MockHttpServletRequest httpServletRequest = new MockHttpServletRequest();
 %0A%0A      

Use spring mock request and response in unit tests