package com.yahoo.omid.tso;

import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyLong;
import static org.mockito.Matchers.anyString;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.verify;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;

import java.io.IOException;

import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

import com.yahoo.omid.metrics.MetricsRegistry;
import com.yahoo.omid.timestamp.storage.TimestampStorage;

public class TestTimestampOracle {

    private static final Logger LOG = LoggerFactory.getLogger(TestTimestampOracle.class);

    @Mock
    private MetricsRegistry metrics;
    @Mock
    private Panicker panicker;
    @Mock
    private TimestampStorage timestampStorage;

    // Component under test
    @InjectMocks
    private TimestampOracleImpl tso;

    @BeforeMethod(alwaysRun = true)
    public void initMocksAndComponents() {

        MockitoAnnotations.initMocks(this);

    }

    @Test
    public void testMonotonicTimestampGrowth() throws Exception {

        // Intialize component under test
        tso.initialize();

        long last = tso.next();
        for (int i = 0; i < (3 * TimestampOracleImpl.TIMESTAMP_BATCH); i++) {
            long current = tso.next();
            assertEquals(current, last + 1, "Not monotonic growth");
            last = current;
        }
        assertTrue(tso.getLast() == last);
        LOG.info("Last timestamp: {}", last);
    }

    @Test
    public void testTimestampOraclePanic() throws Exception {

        // Intialize component under test
        tso.initialize();

        // Cause an exception when updating the max timestamp
        doThrow(new RuntimeException("Out of memory or something"))
            .when(timestampStorage).updateMaxTimestamp(anyLong(), anyLong());

        // Make the previous exception to be thrown
        Thread allocThread = new Thread("AllocThread") {
            @Override
            public void run() {
                try {
                    while (true) {
                        tso.next();
                    }
                } catch (IOException ioe) {
                    LOG.error("Shouldn't occur");
                }
            }
        };
        allocThread.start();

        // Verify that it has blown up
        verify(panicker, timeout(1000).atLeastOnce()).panic(anyString(), any(Throwable.class));
    }
}

@@ -33,378 +33,105 @@
 ort 
-static org.mockito.Matchers.any;%0Aimport static org.mockito.Matchers.anyLong;%0Aimport static org.mockito.Matchers.anyString;%0Aimport static org.mockito.Mockito.doThrow;%0Aimport static org.mockito.Mockito.timeout;%0Aimport static org.mockito.Mockito.verify;%0Aimport static org.testng.Assert.assertEquals;%0Aimport static org.testng.Assert.assertTrue;%0A%0Aimport java.io.IOException;%0A
+com.yahoo.omid.metrics.MetricsRegistry;%0Aimport com.yahoo.omid.timestamp.storage.TimestampStorage;
 %0Aimp
@@ -223,16 +223,100 @@
 ations;%0A
+import org.mockito.invocation.InvocationOnMock;%0Aimport org.mockito.stubbing.Answer;%0A
 import o
@@ -456,103 +456,424 @@
 ort 
-com.yahoo.omid.metrics.MetricsRegistry;%0Aimport com.yahoo.omid.timestamp.storage.TimestampStorag
+java.io.IOException;%0Aimport java.util.concurrent.CountDownLatch;%0A%0Aimport static org.mockito.Matchers.any;%0Aimport static org.mockito.Matchers.anyLong;%0Aimport static org.mockito.Matchers.anyString;%0Aimport static org.mockito.Mockito.atLeastOnce;%0Aimport static org.mockito.Mockito.doAnswer;%0Aimport static org.mockito.Mockito.verify;%0Aimport static org.testng.Assert.assertEquals;%0Aimport static org.testng.Assert.assertTru
 e;%0A%0A
@@ -1223,18 +1223,30 @@
 leImpl t
-so
+imestampOracle
 ;%0A%0A    @
@@ -1373,32 +1373,50 @@
     %7D%0A%0A    @Test
+(timeOut = 10_000)
 %0A    public void
@@ -1510,34 +1510,46 @@
 r test%0A        t
-so
+imestampOracle
 .initialize();%0A%0A
@@ -1565,26 +1565,38 @@
 ong last = t
-so
+imestampOracle
 .next();%0A   
@@ -1698,18 +1698,30 @@
 rent = t
-so
+imestampOracle
 .next();
@@ -1848,18 +1848,30 @@
 rtTrue(t
-so
+imestampOracle
 .getLast
@@ -1945,16 +1945,34 @@
    @Test
+(timeOut = 10_000)
 %0A    pub
@@ -2004,16 +2004,42 @@
 clePanic
+sWhenTheStorageHasProblems
 () throw
@@ -2104,18 +2104,30 @@
        t
-so
+imestampOracle
 .initial
@@ -2208,16 +2208,308 @@
     
-doThrow(
+final CountDownLatch updateMaxTimestampMethodCalled = new CountDownLatch(1);%0A        doAnswer(new Answer() %7B%0A            @Override%0A            public Object answer(InvocationOnMock invocationOnMock) throws Throwable %7B%0A                updateMaxTimestampMethodCalled.countDown();%0A                throw 
 new 
@@ -2554,17 +2554,17 @@
 ething%22)
-)
+;
 %0A       
@@ -2568,16 +2568,28 @@
         
+%7D%0A        %7D)
 .when(ti
@@ -2892,18 +2892,30 @@
        t
-so
+imestampOracle
 .next();
@@ -3100,24 +3100,73 @@
 d.start();%0A%0A
+        updateMaxTimestampMethodCalled.await();%0A%0A
         // V
@@ -3221,22 +3221,8 @@
 er, 
-timeout(1000).
 atLe
@@ -3280,10 +3280,11 @@
 ;%0A    %7D%0A
+%0A
 %7D%0A

Add countdown latch to avoid errors in test verifications