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

import com.github.bedrin.jdbc.sniffer.Sniffer;
import com.github.bedrin.jdbc.sniffer.Spy;
import com.github.bedrin.jdbc.sniffer.WrongNumberOfQueriesError;
import com.github.bedrin.jdbc.sniffer.Expectation;
import com.github.bedrin.jdbc.sniffer.Expectations;
import com.github.bedrin.jdbc.sniffer.NoQueriesAllowed;
import com.github.bedrin.jdbc.sniffer.util.ExceptionUtil;
import org.springframework.test.context.TestContext;
import org.springframework.test.context.support.AbstractTestExecutionListener;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * @since 2.2
 */
public class QueryCounterListener extends AbstractTestExecutionListener {

    private static final String SPY_ATTRIBUTE_NAME = "spy";

    @Override
    public void beforeTestMethod(TestContext testContext) throws Exception {

        Method testMethod = testContext.getTestMethod();

        Expectations expectations = testMethod.getAnnotation(Expectations.class);
        Expectation expectation = testMethod.getAnnotation(Expectation.class);
        NoQueriesAllowed notAllowedQueries = testMethod.getAnnotation(NoQueriesAllowed.class);

        // If no annotations present, check the test class and its superclasses
        for (Class<?> testClass = testMethod.getDeclaringClass();
             null == expectations && null == expectation && null == notAllowedQueries && !Object.class.equals(testClass);
             testClass = testClass.getSuperclass()) {
            expectations = testClass.getAnnotation(Expectations.class);
            expectation = testClass.getAnnotation(Expectation.class);
            notAllowedQueries = testClass.getAnnotation(NoQueriesAllowed.class);
        }

        if (null != expectation && null != notAllowedQueries) {
            throw new IllegalArgumentException("Cannot specify @Expectation and @NotAllowedQueries on one test method");
        } else if (null != expectations && null != notAllowedQueries) {
            throw new IllegalArgumentException("Cannot specify @Expectations and @NotAllowedQueries on one test method");
        } else if (null != expectations || null != expectation) {

            List<Expectation> expectationList = new ArrayList<Expectation>();

            if (null != expectation) {
                expectationList.add(expectation);
            }

            if (null != expectations) {
                expectationList.addAll(Arrays.asList(expectations.value()));
            }

            for (Expectation expectation1 : expectationList) {
                if (expectation1.value() != -1) {
                    if (expectation1.atMost() != -1 || expectation1.atLeast() != -1) {
                        throw new IllegalArgumentException("Cannot specify value parameter together with atLeast or atMost parameters");
                    }
                }
            }

            testContext.setAttribute(SPY_ATTRIBUTE_NAME,
                    Sniffer.expect(expectationList)
            );

        } else if (null != notAllowedQueries) {
            testContext.setAttribute(SPY_ATTRIBUTE_NAME,
                    Sniffer.expect(Collections.singletonList(NoQueriesAllowed.class.getAnnotation(Expectation.class)))
            );
        }

    }

    @Override
    public void afterTestMethod(TestContext testContext) throws Exception {

        Object spyAttribute = testContext.getAttribute(SPY_ATTRIBUTE_NAME);

        if (null != spyAttribute) {

            Spy spy = (Spy) spyAttribute;

            try {
                spy.close();
            } catch (WrongNumberOfQueriesError jdbcSnifferError) {

                Throwable throwable = testContext.getTestException();
                if (null != throwable) {
                    if (!ExceptionUtil.addSuppressed(throwable, jdbcSnifferError)) {
                        jdbcSnifferError.printStackTrace();
                    }
                } else {
                    throw jdbcSnifferError;
                }

            }

        }

    }
}

@@ -546,16 +546,68 @@
 tener;%0A%0A
+import java.lang.reflect.InvocationTargetException;%0A
 import j
@@ -896,16 +896,1497 @@
 %22spy%22;%0A%0A
+    private static final Method GET_TEST_METHOD;%0A    private static final Method SET_ATTRIBUTE_METHOD;%0A    private static final Method GET_ATTRIBUTE_METHOD;%0A%0A    static %7B%0A        Method getTestMethod = null;%0A        Method setAttributeMethod = null;%0A        Method getAttributeMethod = null;%0A        try %7B%0A            getTestMethod = TestContext.class.getMethod(%22getTestMethod%22);%0A            setAttributeMethod = TestContext.class.getMethod(%22setAttribute%22, String.class, Object.class);%0A            getAttributeMethod = TestContext.class.getMethod(%22getAttribute%22, String.class);%0A        %7D catch (NoSuchMethodException e) %7B%0A            e.printStackTrace(); // TODO: what to do with the exception?%0A        %7D%0A        GET_TEST_METHOD = getTestMethod;%0A        SET_ATTRIBUTE_METHOD = setAttributeMethod;%0A        GET_ATTRIBUTE_METHOD = getAttributeMethod;%0A    %7D%0A%0A    private static Method getTestMethod(TestContext testContext) throws InvocationTargetException, IllegalAccessException %7B%0A        return Method.class.cast(GET_TEST_METHOD.invoke(testContext));%0A    %7D%0A%0A    private static void setAttribute(TestContext testContext, String name, Object value) throws InvocationTargetException, IllegalAccessException %7B%0A        SET_ATTRIBUTE_METHOD.invoke(testContext, name, value);%0A    %7D%0A%0A    private static Object getAttribute(TestContext testContext, String name) throws InvocationTargetException, IllegalAccessException %7B%0A        return GET_ATTRIBUTE_METHOD.invoke(testContext, name);%0A    %7D%0A%0A
     @Ove
@@ -2497,28 +2497,16 @@
 ethod = 
-testContext.
 getTestM
@@ -2511,16 +2511,27 @@
 tMethod(
+testContext
 );%0A%0A    
@@ -4500,36 +4500,24 @@
             
-testContext.
 setAttribute
@@ -4509,32 +4509,45 @@
    setAttribute(
+testContext, 
 SPY_ATTRIBUTE_NA
@@ -4549,36 +4549,16 @@
 TE_NAME,
-%0A                   
  Sniffer
@@ -4581,29 +4581,16 @@
 ionList)
-%0A            
 );%0A%0A    
@@ -4645,28 +4645,16 @@
         
-testContext.
 setAttri
@@ -4650,32 +4650,45 @@
    setAttribute(
+testContext, 
 SPY_ATTRIBUTE_NA
@@ -4964,28 +4964,16 @@
 ibute = 
-testContext.
 getAttri
@@ -4977,16 +4977,29 @@
 tribute(
+testContext, 
 SPY_ATTR

Call spring TestContext via reflection to ensure compatibility with multiple versions