package com.indeed.proctor.common;

import com.indeed.proctor.common.model.TestMatrixArtifact;
import com.indeed.util.core.DataLoadTimer;
import org.apache.el.lang.FunctionMapperImpl;
import org.easymock.classextension.EasyMock;
import org.junit.Test;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;

public class AbstractProctorLoaderTest {

    private final DataLoadTimer dataLoaderTimerMock = EasyMock.createMock(DataLoadTimer.class);
    private final TestProctorLoader loader = createTestProctorLoader(dataLoaderTimerMock);

    @Test
    public void testLoaderTimerFunctionsNoLoad() {
        EasyMock.expect(dataLoaderTimerMock.getSecondsSinceLastLoadCheck())
                .andReturn(null)
                .once();
        EasyMock.expect(dataLoaderTimerMock.isLoadedDataSuccessfullyRecently())
                .andReturn(false)
                .once();
        EasyMock.replay(dataLoaderTimerMock);

        assertNull(loader.secondsSinceLastReloadAttempt());
        assertFalse(loader.isLoadedDataSuccessfullyRecently());
    }

    @Test
    public void testLoaderTimerFunctionsLoadSuccess() {
        EasyMock.expect(dataLoaderTimerMock.isLoadedDataSuccessfullyRecently())
                .andReturn(true)
                .once();
        EasyMock.expect(dataLoaderTimerMock.getSecondsSinceLastLoadCheck())
                .andReturn(42)
                .once();
        EasyMock.replay(dataLoaderTimerMock);

        assertEquals(42, (int) loader.secondsSinceLastReloadAttempt());
        assertTrue(loader.isLoadedDataSuccessfullyRecently());
    }

    private static TestProctorLoader createTestProctorLoader(final DataLoadTimer dataLoaderTimer) {
        return new TestProctorLoader(dataLoaderTimer);
    }

    private static class TestProctorLoader extends AbstractProctorLoader {

        public TestProctorLoader(final DataLoadTimer dataLoaderTimer) {
            super(TestProctorLoader.class, new ProctorSpecification(), new FunctionMapperImpl());
            this.dataLoadTimer = dataLoaderTimer;
        }

        @Nullable
        @Override
        TestMatrixArtifact loadTestMatrix() {
            return null;
        }

        @Nonnull
        @Override
        String getSource() {
            return "dummy";
        }
    }
}
@@ -222,16 +222,41 @@
 syMock;%0A
+import org.junit.Before;%0A
 import o
@@ -569,35 +569,135 @@
 private 
-final D
+DataLoadTimer d
 ataLoad
+er
 Timer
+Mock;%0A    private TestProctorLoader loader;%0A%0A    @Before%0A    public void setUp() %7B%0A       
  dataLoa
@@ -761,39 +761,11 @@
     
-private final TestProctorLoader
+   
  loa
@@ -811,24 +811,30 @@
 rTimerMock);
+%0A    %7D
 %0A%0A    @Test%0A

Use @Before in test