package alluxio.worker.block;

import alluxio.ConfigurationRule;
import alluxio.PropertyKey;
import alluxio.underfs.SeekableUnderFileInputStream;
import alluxio.underfs.UnderFileSystem;
import alluxio.underfs.options.OpenOptions;

import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;

import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

public class UfsInputStreamManagerTest {
  private static final String FILE_NAME = "/test";

  private UnderFileSystem mUfs;
  private List<SeekableUnderFileInputStream> mSeekableInStreams;
  private UfsInputStreamManager mManager;

  public ConfigurationRule mConfigurationRule =
      new ConfigurationRule(new HashMap<PropertyKey, String>() {
        {
          put(PropertyKey.WORKER_UFS_INSTREAM_CACHE_EXPIRE_MS, "200000");
        }
      });

  @Before
  public void before() throws Exception {
    mSeekableInStreams = new ArrayList<>();
    mUfs = Mockito.mock(UnderFileSystem.class);
    Mockito.when(mUfs.isSeekable()).thenReturn(true);
    for(int i=0;i<3;i++) {
      SeekableUnderFileInputStream instream=createMockedSeekableStream();
      Mockito.when(mUfs.open(Mockito.eq(FILE_NAME), Mockito.any(OpenOptions.class))).thenReturn(instream);
      mSeekableInStreams.add(instream);
    }
    mManager = new UfsInputStreamManager();
  }

  private SeekableUnderFileInputStream createMockedSeekableStream() {
    SeekableUnderFileInputStream instream = Mockito.mock(SeekableUnderFileInputStream.class);
    Mockito.when(instream.getFilePath()).thenReturn(FILE_NAME);
    Mockito.doCallRealMethod().when(instream).setResourceId(Mockito.any(Long.class));
    Mockito.when(instream.getResourceId()).thenCallRealMethod();
    return instream;
  }

  @Test
  public void testCheckInAndOut() throws Exception {
    mSeekableInStreams = new ArrayList<>();
    SeekableUnderFileInputStream mockedStrem = createMockedSeekableStream();
    Mockito.when(mUfs.open(Mockito.eq(FILE_NAME), Mockito.any(OpenOptions.class))).thenReturn(mockedStrem)
        .thenThrow(new IllegalStateException("Should only be called once"));
    mSeekableInStreams.add(mockedStrem);

    // check out a stream
    InputStream instream1 = mManager.checkOut(mUfs, FILE_NAME, 2);
    // check in back
    mManager.checkIn(instream1);
 // check out a stream again
    InputStream instream2 = mManager.checkOut(mUfs, FILE_NAME, 4);

    Assert.assertEquals(mockedStrem, instream1);
    Assert.assertEquals(mockedStrem, instream2);
    // ensure the second time the checked out instream is the same one but repositioned
    Mockito.verify(mockedStrem).seek(4);
  }
}

@@ -326,16 +326,42 @@
 ckito;%0A%0A
+import java.io.Closeable;%0A
 import j
@@ -695,226 +695,8 @@
 r;%0A%0A
-  public ConfigurationRule mConfigurationRule =%0A      new ConfigurationRule(new HashMap%3CPropertyKey, String%3E() %7B%0A        %7B%0A          put(PropertyKey.WORKER_UFS_INSTREAM_CACHE_EXPIRE_MS, %22200000%22);%0A        %7D%0A      %7D);%0A%0A
   @B
@@ -988,24 +988,68 @@
 Stream();%0A  
+    mSeekableInStreams.add(instream);%0A    %7D%0A
     Mockito.
@@ -1130,27 +1130,19 @@
 nReturn(
-instream);%0A
+%0A  
       mS
@@ -1158,36 +1158,78 @@
 Streams.
-add(instream);%0A    %7D
+get(0), mSeekableInStreams.get(1), mSeekableInStreams.get(2));
 %0A    mMa
@@ -2230,16 +2230,19 @@
 ream1);%0A
+   
  // chec
@@ -2560,11 +2560,1125 @@
 (4);%0A  %7D
+%0A%0A  @Test%0A  public void testMultipleCheckIn() throws Exception %7B%0A    mManager.checkOut(mUfs, FILE_NAME, 2);%0A    mManager.checkOut(mUfs, FILE_NAME, 4);%0A    mManager.checkOut(mUfs, FILE_NAME, 6);%0A    // 3 different input streams are checked out%0A    Mockito.verify(mUfs, Mockito.times(3)).open(Mockito.eq(FILE_NAME),%0A        Mockito.any(OpenOptions.class));%0A  %7D%0A%0A  @Test%0A  public void testExpire() throws Exception %7B%0A    try (Closeable r = new ConfigurationRule(new HashMap%3CPropertyKey, String%3E() %7B%0A      %7B%0A        put(PropertyKey.WORKER_UFS_INSTREAM_CACHE_EXPIRE_MS, %222%22);%0A      %7D%0A    %7D).toResource()) %7B%0A      mManager = new UfsInputStreamManager();%0A      InputStream instream;%0A      // check out a stream%0A      instream = mManager.checkOut(mUfs, FILE_NAME, 2);%0A      // check in back%0A      mManager.checkIn(instream);%0A      Thread.sleep(10);%0A      // check out another stream should trigger the timeout%0A      mManager.checkOut(mUfs, FILE_NAME, 4);%0A      Mockito.verify(mSeekableInStreams.get(0), Mockito.timeout(1000000).times(1)).close();%0A    %7D;%0A  %7D%0A%0A  @Test%0A  public void testConcurrency() throws Exception %7B%0A%0A  %7D
 %0A%7D%0A

add more tests