text:
package org.apache.jackrabbit.oak.plugins.index.lucene;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

import javax.management.openmbean.TabularData;

import com.google.common.base.StandardSystemProperty;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.google.common.util.concurrent.ForwardingListeningExecutorService;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import org.apache.commons.io.FileUtils;
import org.apache.jackrabbit.oak.commons.IOUtils;
import org.apache.jackrabbit.oak.spi.state.NodeBuilder;
import org.apache.jackrabbit.oak.spi.state.NodeState;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FilterDirectory;
import org.apache.lucene.store.IOContext;
import org.apache.lucene.store.IndexInput;
import org.apache.lucene.store.IndexOutput;
import org.apache.lucene.store.RAMDirectory;
import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import static com.google.common.collect.Lists.newArrayList;
import static com.google.common.collect.Sets.newHashSet;
import static com.google.common.util.concurrent.MoreExecutors.sameThreadExecutor;
import static org.apache.jackrabbit.oak.plugins.index.IndexConstants.REINDEX_COUNT;
import static org.apache.jackrabbit.oak.plugins.nodetype.write.InitialContent.INITIAL_CONTENT;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.junit.Assume.assumeTrue;

public class IndexCopierTest {
    private Random rnd = new Random();
    private int maxFileSize = 7896;

    private NodeState root = INITIAL_CONTENT;

    @Rule
    public TemporaryFolder temporaryFolder = new TemporaryFolder();

    private NodeBuilder builder = root.builder();

    @Test
    public void basicTest() throws Exception{
        Directory baseDir = new RAMDirectory();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier c1 = new RAMIndexCopier(baseDir, sameThreadExecutor(), getWorkDir());

        Directory remote = new RAMDirectory();
        Directory wrapped = c1.wrapForRead("/foo", defn, remote);

        byte[] t1 = writeFile(remote , "t1");
        byte[] t2 = writeFile(remote , "t2");

        assertEquals(2, wrapped.listAll().length);

        assertTrue(wrapped.fileExists("t1"));
        assertTrue(wrapped.fileExists("t2"));

        assertEquals(t1.length, wrapped.fileLength("t1"));
        assertEquals(t2.length, wrapped.fileLength("t2"));

        readAndAssert(wrapped, "t1", t1);

        //t1 should now be added to testDir
        assertTrue(baseDir.fileExists("t1"));
    }

    @Test
    public void basicTestWithPrefetch() throws Exception{
        Directory baseDir = new RAMDirectory();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier c1 = new RAMIndexCopier(baseDir, sameThreadExecutor(), getWorkDir(), true);

        Directory remote = new RAMDirectory();

        byte[] t1 = writeFile(remote, "t1");
        byte[] t2 = writeFile(remote , "t2");

        Directory wrapped = c1.wrapForRead("/foo", defn, remote);
        assertEquals(2, wrapped.listAll().length);

        assertTrue(wrapped.fileExists("t1"));
        assertTrue(wrapped.fileExists("t2"));

        assertTrue(baseDir.fileExists("t1"));
        assertTrue(baseDir.fileExists("t2"));

        assertEquals(t1.length, wrapped.fileLength("t1"));
        assertEquals(t2.length, wrapped.fileLength("t2"));

        readAndAssert(wrapped, "t1", t1);

    }

    @Test
    public void basicTestWithFS() throws Exception{
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier c1 = new IndexCopier(sameThreadExecutor(), getWorkDir());

        Directory remote = new RAMDirectory();
        Directory wrapped = c1.wrapForRead("/foo", defn, remote);

        byte[] t1 = writeFile(remote, "t1");
        byte[] t2 = writeFile(remote , "t2");

        assertEquals(2, wrapped.listAll().length);

        assertTrue(wrapped.fileExists("t1"));
        assertTrue(wrapped.fileExists("t2"));

        assertEquals(t1.length, wrapped.fileLength("t1"));
        assertEquals(t2.length, wrapped.fileLength("t2"));

        readAndAssert(wrapped, "t1", t1);

        //t1 should now be added to testDir
        File indexBaseDir = c1.getIndexDir("/foo");
        File indexDir = new File(indexBaseDir, "0");
        assertTrue(new File(indexDir, "t1").exists());

        TabularData td = c1.getIndexPathMapping();
        assertEquals(1, td.size());
    }

    @Test
    public void deleteOldPostReindex() throws Exception{
        assumeNotWindows();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier c1 = new IndexCopier(sameThreadExecutor(), getWorkDir());

        Directory remote = new CloseSafeDir();
        Directory w1 = c1.wrapForRead("/foo", defn, remote);

        byte[] t1 = writeFile(remote, "t1");
        byte[] t2 = writeFile(remote , "t2");

        readAndAssert(w1, "t1", t1);
        readAndAssert(w1, "t2", t2);

        //t1 should now be added to testDir
        File indexBaseDir = c1.getIndexDir("/foo");
        File indexDir = new File(indexBaseDir, "0");
        assertTrue(new File(indexDir, "t1").exists());

        builder.setProperty(REINDEX_COUNT, 1);
        defn = new IndexDefinition(root, builder.getNodeState());

        //Close old version
        w1.close();
        //Get a new one with updated reindexCount
        Directory w2 = c1.wrapForRead("/foo", defn, remote);

        readAndAssert(w2, "t1", t1);

        w2.close();
        assertFalse("Old index directory should have been removed", indexDir.exists());

        File indexDir2 = new File(indexBaseDir, "1");
        assertTrue(new File(indexDir2, "t1").exists());
    }

    @Test
    public void concurrentRead() throws Exception{
        Directory baseDir = new RAMDirectory();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        CollectingExecutor executor = new CollectingExecutor();

        IndexCopier c1 = new RAMIndexCopier(baseDir, executor, getWorkDir());

        TestRAMDirectory remote = new TestRAMDirectory();
        Directory wrapped = c1.wrapForRead("/foo", defn, remote);

        byte[] t1 = writeFile(remote , "t1");

        //1. Trigger a read which should go to remote
        readAndAssert(wrapped, "t1", t1);
        assertEquals(1, c1.getScheduledForCopyCount());
        assertEquals(1, remote.openedFiles.size());
        assertEquals(1, executor.commands.size());

        //2. Trigger another read and this should also be
        //served from remote
        readAndAssert(wrapped, "t1", t1);
        assertEquals(1, c1.getScheduledForCopyCount());
        assertEquals(2, remote.openedFiles.size());
        //Second read should not add a new copy task
        assertEquals(1, executor.commands.size());

        //3. Perform copy
        executor.executeAll();
        remote.reset();

        //4. Now read again after copy is done
        readAndAssert(wrapped, "t1", t1);
        // Now read should be served from local and not from remote
        assertEquals(0, remote.openedFiles.size());
        assertEquals(0, c1.getScheduledForCopyCount());
    }

    @Test
    public void copyInProgressStats() throws Exception{
        Directory baseDir = new RAMDirectory();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());

        final List<ListenableFuture<?>> submittedTasks = Lists.newArrayList();
        ExecutorService executor = new ForwardingListeningExecutorService() {
            @Override
            protected ListeningExecutorService delegate() {
                return MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor());
            }

            @Override
            public void execute(Runnable command) {
                submittedTasks.add(super.submit(command));
            }
        };

        IndexCopier c1 = new RAMIndexCopier(baseDir, executor, getWorkDir());

        final CountDownLatch copyProceed = new CountDownLatch(1);
        final CountDownLatch copyRequestArrived = new CountDownLatch(1);
        TestRAMDirectory remote = new TestRAMDirectory(){
            @Override
            public void copy(Directory to, String src, String dest, IOContext context) throws IOException {
                copyRequestArrived.countDown();
                try {
                    copyProceed.await();
                } catch (InterruptedException e) {

                }
                super.copy(to, src, dest, context);
            }
        };
        Directory wrapped = c1.wrapForRead("/foo", defn, remote);

        byte[] t1 = writeFile(remote , "t1");

        //1. Trigger a read which should go to remote
        readAndAssert(wrapped, "t1", t1);
        copyRequestArrived.await();
        assertEquals(1, c1.getCopyInProgressCount());
        assertEquals(1, remote.openedFiles.size());

        //2. Trigger another read and this should also be
        //served from remote
        readAndAssert(wrapped, "t1", t1);
        assertEquals(1, c1.getCopyInProgressCount());
        assertEquals(IOUtils.humanReadableByteCount(t1.length), c1.getCopyInProgressSize());
        assertEquals(1, c1.getCopyInProgressDetails().length);
        System.out.println(Arrays.toString(c1.getCopyInProgressDetails()));
        assertEquals(2, remote.openedFiles.size());

        //3. Perform copy
        copyProceed.countDown();
        Futures.allAsList(submittedTasks).get();
        remote.reset();

        //4. Now read again after copy is done
        readAndAssert(wrapped, "t1", t1);
        // Now read should be served from local and not from remote
        assertEquals(0, remote.openedFiles.size());
        assertEquals(0, c1.getCopyInProgressCount());

        executor.shutdown();
    }

    /**
     * Test for the case where local directory is opened already contains
     * the index files and in such a case file should not be read from remote
     */
    @Test
    public void reuseLocalDir() throws Exception{
        Directory baseDir = new RAMDirectory();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier c1 = new RAMIndexCopier(baseDir, sameThreadExecutor(), getWorkDir());

        TestRAMDirectory remote = new TestRAMDirectory();
        Directory wrapped = c1.wrapForRead("/foo", defn, remote);

        byte[] t1 = writeFile(remote, "t1");

        //1. Read for the first time should be served from remote
        readAndAssert(wrapped, "t1", t1);
        assertEquals(1, remote.openedFiles.size());

        //2. Reuse the testDir and read again
        Directory wrapped2 = c1.wrapForRead("/foo", defn, remote);
        remote.reset();

        //3. Now read should be served from local
        readAndAssert(wrapped2, "t1", t1);
        assertEquals(0, remote.openedFiles.size());

        //Now check if local file gets corrupted then read from remote
        Directory wrapped3 = c1.wrapForRead("/foo", defn, remote);
        remote.reset();

        //4. Corrupt the local copy
        writeFile(baseDir, "t1");

        //Now read would be done from remote
        readAndAssert(wrapped3, "t1", t1);
        assertEquals(1, remote.openedFiles.size());
    }

    @Test
    public void deleteCorruptedFile() throws Exception{
        Directory baseDir = new RAMDirectory();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        RAMIndexCopier c1 = new RAMIndexCopier(baseDir, sameThreadExecutor(), getWorkDir());

        Directory remote = new RAMDirectory(){
            @Override
            public IndexInput openInput(String name, IOContext context) throws IOException {
                throw new IllegalStateException("boom");
            }
        };

        String fileName = "failed.txt";
        Directory wrapped = c1.wrapForRead("/foo", defn, remote);

        byte[] t1 = writeFile(remote , fileName);

        try {
            readAndAssert(wrapped, fileName, t1);
            fail("Read of file should have failed");
        } catch (IllegalStateException ignore){

        }

        assertFalse(c1.baseDir.fileExists(fileName));
    }

    @Test
    public void deletesOnClose() throws Exception{
        //Use a close safe dir. In actual case the FSDir would
        //be opened on same file system hence it can retain memory
        //but RAMDirectory does not retain memory hence we simulate
        //that by not closing the RAMDir and reuse it
        Directory baseDir = new CloseSafeDir();

        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier c1 = new RAMIndexCopier(baseDir, sameThreadExecutor(), getWorkDir());

        Directory r1 = new RAMDirectory();

        byte[] t1 = writeFile(r1, "t1");
        byte[] t2 = writeFile(r1 , "t2");

        Directory w1 = c1.wrapForRead("/foo", defn, r1);
        readAndAssert(w1, "t1", t1);
        readAndAssert(w1, "t2", t2);

        // t1 and t2 should now be present in local (base dir which back local)
        assertTrue(baseDir.fileExists("t1"));
        assertTrue(baseDir.fileExists("t2"));

        Directory r2 = new RAMDirectory();
        copy(r1, r2);
        r2.deleteFile("t1");

        Directory w2 = c1.wrapForRead("/foo", defn, r2);

        //Close would trigger removal of file which are not present in remote
        w2.close();

        assertFalse("t1 should have been deleted", baseDir.fileExists("t1"));
        assertTrue(baseDir.fileExists("t2"));
    }

    @Test
    public void failureInDelete() throws Exception{
        final Set<String> testFiles = new HashSet<String>();
        Directory baseDir = new CloseSafeDir() {
            @Override
            public void deleteFile(String name) throws IOException {
                if (testFiles.contains(name)){
                    throw new IOException("Not allowed to delete " + name);
                }
                super.deleteFile(name);
            }
        };

        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier c1 = new RAMIndexCopier(baseDir, sameThreadExecutor(), getWorkDir());

        Directory r1 = new RAMDirectory();

        byte[] t1 = writeFile(r1, "t1");
        byte[] t2 = writeFile(r1 , "t2");

        Directory w1 = c1.wrapForRead("/foo", defn, r1);
        readAndAssert(w1, "t1", t1);
        readAndAssert(w1, "t2", t2);

        // t1 and t2 should now be present in local (base dir which back local)
        assertTrue(baseDir.fileExists("t1"));
        assertTrue(baseDir.fileExists("t2"));

        Directory r2 = new CloseSafeDir();
        copy(r1, r2);
        r2.deleteFile("t1");

        Directory w2 = c1.wrapForRead("/foo", defn, r2);

        //Close would trigger removal of file which are not present in remote
        testFiles.add("t1");
        w2.close();

        assertEquals(1, c1.getFailedToDeleteFiles().size());
        IndexCopier.LocalIndexFile testFile = c1.getFailedToDeleteFiles().values().iterator().next();

        assertEquals(1, testFile.getDeleteAttemptCount());
        assertEquals(IOUtils.humanReadableByteCount(t1.length), c1.getGarbageSize());
        assertEquals(1, c1.getGarbageDetails().length);

        Directory w3 = c1.wrapForRead("/foo", defn, r2);
        w3.close();
        assertEquals(2, testFile.getDeleteAttemptCount());

        //Now let the file to be deleted
        testFiles.clear();

        Directory w4 = c1.wrapForRead("/foo", defn, r2);
        w4.close();

        //No pending deletes left
        assertEquals(0, c1.getFailedToDeleteFiles().size());
    }

    @Test
    public void deletedOnlyFilesForOlderVersion() throws Exception{
        Directory baseDir = new CloseSafeDir();

        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier copier = new RAMIndexCopier(baseDir, sameThreadExecutor(), getWorkDir());

        //1. Open a local and read t1 from remote
        Directory remote1 = new RAMDirectory();
        byte[] t1 = writeFile(remote1, "t1");

        Directory local1 = copier.wrapForRead("/foo", defn, remote1);
        readAndAssert(local1, "t1", t1);

        //While local1 is open , open another local2 and read t2
        Directory remote2 = new RAMDirectory();
        byte[] t2 = writeFile(remote2, "t2");

        Directory local2 = copier.wrapForRead("/foo", defn, remote2);
        readAndAssert(local2, "t2", t2);

        //Close local1
        local1.close();

        //t2 should still be readable
        readAndAssert(local2, "t2", t2);
    }

    @Test
    public void wrapForWriteWithoutIndexPath() throws Exception{
        assumeNotWindows();
        Directory remote = new CloseSafeDir();

        IndexCopier copier = new IndexCopier(sameThreadExecutor(), getWorkDir());

        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        Directory dir = copier.wrapForWrite(defn, remote, false);

        byte[] t1 = writeFile(dir, "t1");

        dir.close();

        readAndAssert(remote, "t1", t1);
        //Work dir must be empty post close
        assertArrayEquals(FileUtils.EMPTY_FILE_ARRAY, copier.getIndexWorkDir().listFiles());
    }

    @Test
    public void wrapForWriteWithIndexPath() throws Exception{
        assumeNotWindows();
        Directory remote = new CloseSafeDir();

        IndexCopier copier = new IndexCopier(sameThreadExecutor(), getWorkDir());

        builder.setProperty(LuceneIndexConstants.INDEX_PATH, "foo");
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        Directory dir = copier.wrapForWrite(defn, remote, false);

        byte[] t1 = writeFile(dir, "t1");

        dir.close();

        readAndAssert(remote, "t1", t1);
        //Work dir must be empty post close
        List<File> files = new ArrayList<File>(FileUtils.listFiles(copier.getIndexRootDir(), null, true));
        assertEquals(1, files.size());
        assertEquals("t1", files.get(0).getName());
    }

    @Test
    public void copyOnWriteBasics() throws Exception{
        Directory baseDir = new CloseSafeDir();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier copier = new RAMIndexCopier(baseDir, sameThreadExecutor(), getWorkDir());

        Directory remote = new RAMDirectory();
        byte[] t1 = writeFile(remote, "t1");

        //State of remote directory should set before wrapping as later
        //additions would not be picked up given COW assume remote directory
        //to be read only
        Directory local = copier.wrapForWrite(defn, remote, false);

        assertEquals(newHashSet("t1"), newHashSet(local.listAll()));
        assertEquals(t1.length, local.fileLength("t1"));

        byte[] t2 = writeFile(local, "t2");
        assertEquals(newHashSet("t1", "t2"), newHashSet(local.listAll()));
        assertEquals(t2.length, local.fileLength("t2"));

        assertTrue(local.fileExists("t1"));
        assertTrue(local.fileExists("t2"));

        assertTrue("t2 should be copied to remote", remote.fileExists("t2"));

        readAndAssert(local, "t1", t1);
        readAndAssert(local, "t2", t2);

        local.deleteFile("t1");
        assertEquals(newHashSet("t2"), newHashSet(local.listAll()));

        local.deleteFile("t2");
        assertEquals(newHashSet(), newHashSet(local.listAll()));

        try {
            local.fileLength("nonExistentFile");
            fail();
        } catch (FileNotFoundException ignore) {

        }

        try {
            local.openInput("nonExistentFile", IOContext.DEFAULT);
            fail();
        } catch (FileNotFoundException ignore) {

        }

        local.close();
        assertFalse(baseDir.fileExists("t2"));
    }

    /**
     * Checks for the case where if the file exist local before writer starts
     * then those files do not get deleted even if deleted by writer via
     * indexing process from 'baseDir' as they might be in use by existing open
     * indexes
     */
    @Test
    public void cowExistingLocalFileNotDeleted() throws Exception{
        Directory baseDir = new CloseSafeDir();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier copier = new RAMIndexCopier(baseDir, sameThreadExecutor(), getWorkDir());

        Directory remote = new CloseSafeDir();
        byte[] t1 = writeFile(remote, "t1");
        byte[] t2 = writeFile(remote, "t2");
        Directory local = copier.wrapForWrite(defn, remote, false);
        assertEquals(newHashSet("t1", "t2"), newHashSet(local.listAll()));

        byte[] t3 = writeFile(local, "t3");

        //Now pull in the file t1 via CopyOnRead in baseDir
        Directory localForRead = copier.wrapForRead("/foo", defn, remote);
        readAndAssert(localForRead, "t1", t1);

        //File which was copied from remote should not be deleted from baseDir
        //upon delete from local
        assertTrue(baseDir.fileExists("t1"));
        local.deleteFile("t1");
        assertFalse("t1 should be deleted from remote", remote.fileExists("t1"));
        assertFalse("t1 should be deleted from 'local' view also", local.fileExists("t1"));
        assertTrue("t1 should not be deleted from baseDir", baseDir.fileExists("t1"));

        //File which was created only via local SHOULD get removed from
        //baseDir only upon close
        assertTrue(baseDir.fileExists("t3"));
        local.deleteFile("t3");
        assertFalse("t1 should be deleted from remote", local.fileExists("t3"));
        assertTrue("t1 should NOT be deleted from remote", baseDir.fileExists("t3"));

        local.close();
        assertFalse("t3 should also be deleted from local", baseDir.fileExists("t3"));
    }

    @Test
    public void cowReadDoneFromLocalIfFileExist() throws Exception{
        final Set<String> readLocal = newHashSet();
        Directory baseDir = new CloseSafeDir(){
            @Override
            public IndexInput openInput(String name, IOContext context) throws IOException {
                readLocal.add(name);
                return super.openInput(name, context);
            }
        };
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier copier = new RAMIndexCopier(baseDir, sameThreadExecutor(), getWorkDir());

        final Set<String> readRemotes = newHashSet();
        Directory remote = new RAMDirectory() {
            @Override
            public IndexInput openInput(String name, IOContext context) throws IOException {
                readRemotes.add(name);
                return super.openInput(name, context);
            }
        };
        byte[] t1 = writeFile(remote, "t1");
        Directory local = copier.wrapForWrite(defn, remote, false);

        //Read should be served from remote
        readRemotes.clear();readLocal.clear();
        readAndAssert(local, "t1", t1);
        assertEquals(newHashSet("t1"), readRemotes);
        assertEquals(newHashSet(), readLocal);

        //Now pull in the file t1 via CopyOnRead in baseDir
        Directory localForRead = copier.wrapForRead("/foo", defn, remote);
        readAndAssert(localForRead, "t1", t1);

        //Read should be served from local
        readRemotes.clear();readLocal.clear();
        readAndAssert(local, "t1", t1);
        assertEquals(newHashSet(), readRemotes);
        assertEquals(newHashSet("t1"), readLocal);

        local.close();
    }

    @Test
    public void cowCopyDoneOnClose() throws Exception{
        final CollectingExecutor executor = new CollectingExecutor();
        Directory baseDir = new CloseSafeDir();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier copier = new RAMIndexCopier(baseDir, executor, getWorkDir());

        Directory remote = new CloseSafeDir();

        final Directory local = copier.wrapForWrite(defn, remote, false);
        byte[] t1 = writeFile(local, "t1");

        assertTrue(local.fileExists("t1"));
        assertFalse("t1 should NOT be copied to remote", remote.fileExists("t1"));

        //Execute all job
        executor.executeAll();

        assertTrue("t1 should now be copied to remote", remote.fileExists("t1"));

        byte[] t2 = writeFile(local, "t2");
        assertFalse("t2 should NOT be copied to remote", remote.fileExists("t2"));

        final ExecutorService executorService = Executors.newFixedThreadPool(4);
        final CountDownLatch copyLatch = new CountDownLatch(1);
        Future<?> copyTasks = executorService.submit(new Callable<Object>() {
            @Override
            public Object call() throws Exception {
                copyLatch.await();
                //the executor to a proper one as it might happen that
                //STOP task is added post CountingExecutor has executed. Then there
                //would be none to process the STOP. Having a proper executor would
                //handle that case
                executor.setForwardingExecutor(executorService);
                executor.executeAll();
                return null;
            }
        });

        final CountDownLatch closeLatch = new CountDownLatch(1);
        Future<?> closeTasks = executorService.submit(new Callable<Object>() {
            @Override
            public Object call() throws Exception {
                closeLatch.await();
                local.close();
                return null;
            }
        });

        closeLatch.countDown();
        assertFalse("t2 should NOT be copied to remote", remote.fileExists("t2"));

        //Let copy to proceed
        copyLatch.countDown();

        //Now wait for close to finish
        closeTasks.get();
        assertTrue("t2 should now be copied to remote", remote.fileExists("t2"));

        executorService.shutdown();
    }

    @Test
    public void cowCopyDoneOnCloseExceptionHandling() throws Exception{
        final CollectingExecutor executor = new CollectingExecutor();
        Directory baseDir = new CloseSafeDir();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier copier = new RAMIndexCopier(baseDir, executor, getWorkDir());

        Directory remote = new CloseSafeDir();

        final Directory local = copier.wrapForWrite(defn, remote, false);
        byte[] t1 = writeFile(local, "t1");

        assertTrue(local.fileExists("t1"));
        assertFalse("t1 should NOT be copied to remote", remote.fileExists("t1"));

        //Execute all job
        executor.executeAll();

        assertTrue("t1 should now be copied to remote", remote.fileExists("t1"));

        byte[] t2 = writeFile(local, "t2");
        assertFalse("t2 should NOT be copied to remote", remote.fileExists("t2"));

        ExecutorService executorService = Executors.newFixedThreadPool(2);
        final CountDownLatch copyLatch = new CountDownLatch(1);
        Future<?> copyTasks = executorService.submit(new Callable<Object>() {
            @Override
            public Object call() throws Exception {
                copyLatch.await();
                executor.executeAll();
                executor.enableImmediateExecution();
                return null;
            }
        });

        final CountDownLatch closeLatch = new CountDownLatch(1);
        Future<?> closeTasks = executorService.submit(new Callable<Object>() {
            @Override
            public Object call() throws Exception {
                closeLatch.await();
                local.close();
                return null;
            }
        });

        closeLatch.countDown();
        assertFalse("t2 should NOT be copied to remote", remote.fileExists("t2"));

        //Let copy to proceed
        copyLatch.countDown();

        //Now wait for close to finish
        closeTasks.get();
        assertTrue("t2 should now be copied to remote", remote.fileExists("t2"));

        executorService.shutdown();
    }

    @Test
    public void cowFailureInCopy() throws Exception{
        ExecutorService executorService = Executors.newFixedThreadPool(2);
        Directory baseDir = new CloseSafeDir();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier copier = new RAMIndexCopier(baseDir, executorService, getWorkDir());

        final Set<String> toFail = Sets.newHashSet();
        Directory remote = new CloseSafeDir() {
            @Override
            public IndexOutput createOutput(String name, IOContext context) throws IOException {
                if (toFail.contains(name)){
                    throw new RuntimeException("Failing copy for "+name);
                }
                return super.createOutput(name, context);
            }
        };

        final Directory local = copier.wrapForWrite(defn, remote, false);
        toFail.add("t2");
        byte[] t1 = writeFile(local, "t1");
        byte[] t2 = writeFile(local, "t2");

        try {
            local.close();
            fail();
        } catch (IOException ignore){

        }

        executorService.shutdown();
    }

    @Test
    public void cowPoolClosedWithTaskInQueue() throws Exception{
        ExecutorService executorService = Executors.newFixedThreadPool(2);
        Directory baseDir = new CloseSafeDir();
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier copier = new RAMIndexCopier(baseDir, executorService, getWorkDir());

        final Set<String> toPause = Sets.newHashSet();
        final CountDownLatch pauseCopyLatch = new CountDownLatch(1);
        Directory remote = new CloseSafeDir() {
            @Override
            public IndexOutput createOutput(String name, IOContext context) throws IOException {
                if (toPause.contains(name)){
                    try {
                        pauseCopyLatch.await();
                    } catch (InterruptedException ignore) {

                    }
                }
                return super.createOutput(name, context);
            }
        };

        final Directory local = copier.wrapForWrite(defn, remote, false);
        toPause.add("t2");
        byte[] t1 = writeFile(local, "t1");
        byte[] t2 = writeFile(local, "t2");
        byte[] t3 = writeFile(local, "t3");
        byte[] t4 = writeFile(local, "t4");

        final AtomicReference<Throwable> error =
                new AtomicReference<Throwable>();
        Thread closer = new Thread(new Runnable() {
            @Override
            public void run() {
                try {
                    local.close();
                } catch (Throwable e) {
                    e.printStackTrace();
                    error.set(e);
                }
            }
        });

        closer.start();

        copier.close();
        executorService.shutdown();
        executorService.awaitTermination(100, TimeUnit.MILLISECONDS);

        pauseCopyLatch.countDown();
        closer.join();
        assertNotNull("Close should have thrown an exception", error.get());
    }

    /**
     * Test the interaction between COR and COW using same underlying directory
     */
    @Test
    public void cowConcurrentAccess() throws Exception{
        CollectingExecutor executor = new CollectingExecutor();
        ExecutorService executorService = Executors.newFixedThreadPool(2);
        executor.setForwardingExecutor(executorService);

        Directory baseDir = new CloseSafeDir();
        String indexPath = "/foo";
        builder.setProperty(LuceneIndexConstants.INDEX_PATH, indexPath);
        IndexDefinition defn = new IndexDefinition(root, builder.getNodeState());
        IndexCopier copier = new RAMIndexCopier(baseDir, executor, getWorkDir(), true);

        Directory remote = new CloseSafeDir();
        byte[] f1 = writeFile(remote, "f1");

        Directory cor1 = copier.wrapForRead(indexPath, defn, remote);
        readAndAssert(cor1, "f1", f1);
        cor1.close();

        final CountDownLatch pauseCopyLatch = new CountDownLatch(1);
        Directory remote2 = new FilterDirectory(remote) {
            @Override
            public IndexOutput createOutput(String name, IOContext context) throws IOException {
                try {
                    pauseCopyLatch.await();
                } catch (InterruptedException ignore) {

                }
                return super.createOutput(name, context);
            }
        };

        //Start copying a file to remote via COW
        Directory cow1 = copier.wrapForWrite(defn, remote2, false);
        byte[] f2 = writeFile(cow1, "f2");

        //Before copy is done to remote lets delete f1 from remote and
        //open a COR and close it such that it triggers delete of f1
        remote.deleteFile("f1");
        Directory cor2 = copier.wrapForRead(indexPath, defn, remote);

        //Ensure that deletion task submitted to executor get processed immediately
        executor.enableImmediateExecution();
        cor2.close();
        executor.enableDelayedExecution();

        assertFalse(baseDir.fileExists("f1"));
        assertFalse("f2 should not have been copied to remote so far", remote.fileExists("f2"));
        assertTrue("f2 should exist", baseDir.fileExists("f2"));

        pauseCopyLatch.countDown();
        cow1.close();
        assertTrue("f2 should exist", remote.fileExists("f2"));

        executorService.shutdown();
    }

    private byte[] writeFile(Directory dir, String name) throws IOException {
        byte[] data = randomBytes(rnd.nextInt(maxFileSize) + 1);
        IndexOutput o = dir.createOutput(name, IOContext.DEFAULT);
        o.writeBytes(data, data.length);
        o.close();
        return data;
    }

    private byte[] randomBytes(int size) {
        byte[] data = new byte[size];
        rnd.nextBytes(data);
        return data;
    }

    private File getWorkDir(){
        return temporaryFolder.getRoot();
    }

    private static void readAndAssert(Directory wrapped, String fileName, byte[] expectedData) throws IOException {
        IndexInput i = wrapped.openInput(fileName, IOContext.DEFAULT);
        byte[] result = new byte[(int)wrapped.fileLength(fileName)];
        i.readBytes(result, 0, result.length);
        assertTrue(Arrays.equals(expectedData, result));
    }

    private static void copy(Directory source, Directory dest) throws IOException {
        for (String file : source.listAll()) {
            source.copy(dest, file, file, IOContext.DEFAULT);
        }
    }

    private class RAMIndexCopier extends IndexCopier {
        final Directory baseDir;

        public RAMIndexCopier(Directory baseDir, Executor executor, File indexRootDir,
                              boolean prefetchEnabled) throws IOException {
            super(executor, indexRootDir, prefetchEnabled);
            this.baseDir = baseDir;
        }

        public RAMIndexCopier(Directory baseDir, Executor executor, File indexRootDir) throws IOException {
            this(baseDir, executor, indexRootDir, false);
        }

        @Override
        protected Directory createLocalDirForIndexReader(String indexPath, IndexDefinition definition) throws IOException {
            return baseDir;
        }

        @Override
        protected Directory createLocalDirForIndexWriter(IndexDefinition definition) throws IOException {
            return baseDir;
        }
    }

    private static class TestRAMDirectory extends RAMDirectory {
        final List<String> openedFiles = newArrayList();

        @Override
        public IndexInput openInput(String name, IOContext context) throws IOException {
            openedFiles.add(name);
            return super.openInput(name, context);
        }

        public void reset(){
            openedFiles.clear();
        }
    }

    private static class CloseSafeDir extends RAMDirectory {
        @Override
        public void close() {

        }
    }

    private static class CollectingExecutor implements Executor {
        final BlockingQueue<Runnable> commands = new LinkedBlockingQueue<Runnable>();
        private volatile boolean immediateExecution = false;
        private volatile Executor forwardingExecutor;

        @Override
        public void execute(Runnable command) {
            if (immediateExecution){
                command.run();
                return;
            }

            if (forwardingExecutor != null){
                forwardingExecutor.execute(command);
                return;
            }

            commands.add(command);
        }

        void executeAll(){
            Runnable c;
            while ((c = commands.poll()) != null){
                c.run();
            }
        }

        void enableImmediateExecution(){
            immediateExecution = true;
        }

        void enableDelayedExecution(){
            immediateExecution = false;
        }

        void setForwardingExecutor(Executor forwardingExecutor){
            this.forwardingExecutor = forwardingExecutor;
        }
    }

    private static void assumeNotWindows() {
        assumeTrue(!StandardSystemProperty.OS_NAME.value().toLowerCase().contains("windows"));
    }
}


text:
namespace tensorflow {

namespace {

const ::google::protobuf::Descriptor* TensorSliceProto_descriptor_ = NULL;
const ::google::protobuf::internal::GeneratedMessageReflection*
  TensorSliceProto_reflection_ = NULL;
const ::google::protobuf::Descriptor* TensorSliceProto_Extent_descriptor_ = NULL;
const ::google::protobuf::internal::GeneratedMessageReflection*
  TensorSliceProto_Extent_reflection_ = NULL;
struct TensorSliceProto_ExtentOneofInstance {
  ::google::protobuf::int64 length_;
}* TensorSliceProto_Extent_default_oneof_instance_ = NULL;

}  // namespace

void protobuf_AssignDesc_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto() {
  protobuf_AddDesc_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto();
  const ::google::protobuf::FileDescriptor* file =
    ::google::protobuf::DescriptorPool::generated_pool()->FindFileByName(
      "tensorflow/core/framework/tensor_slice.proto");
  GOOGLE_CHECK(file != NULL);
  TensorSliceProto_descriptor_ = file->message_type(0);
  static const int TensorSliceProto_offsets_[1] = {
    GOOGLE_PROTOBUF_GENERATED_MESSAGE_FIELD_OFFSET(TensorSliceProto, extent_),
  };
  TensorSliceProto_reflection_ =
    ::google::protobuf::internal::GeneratedMessageReflection::NewGeneratedMessageReflection(
      TensorSliceProto_descriptor_,
      TensorSliceProto::default_instance_,
      TensorSliceProto_offsets_,
      -1,
      -1,
      -1,
      sizeof(TensorSliceProto),
      GOOGLE_PROTOBUF_GENERATED_MESSAGE_FIELD_OFFSET(TensorSliceProto, _internal_metadata_),
      GOOGLE_PROTOBUF_GENERATED_MESSAGE_FIELD_OFFSET(TensorSliceProto, _is_default_instance_));
  TensorSliceProto_Extent_descriptor_ = TensorSliceProto_descriptor_->nested_type(0);
  static const int TensorSliceProto_Extent_offsets_[3] = {
    GOOGLE_PROTOBUF_GENERATED_MESSAGE_FIELD_OFFSET(TensorSliceProto_Extent, start_),
    PROTO2_GENERATED_DEFAULT_ONEOF_FIELD_OFFSET(TensorSliceProto_Extent_default_oneof_instance_, length_),
    GOOGLE_PROTOBUF_GENERATED_MESSAGE_FIELD_OFFSET(TensorSliceProto_Extent, has_length_),
  };
  TensorSliceProto_Extent_reflection_ =
    ::google::protobuf::internal::GeneratedMessageReflection::NewGeneratedMessageReflection(
      TensorSliceProto_Extent_descriptor_,
      TensorSliceProto_Extent::default_instance_,
      TensorSliceProto_Extent_offsets_,
      -1,
      -1,
      -1,
      TensorSliceProto_Extent_default_oneof_instance_,
      GOOGLE_PROTOBUF_GENERATED_MESSAGE_FIELD_OFFSET(TensorSliceProto_Extent, _oneof_case_[0]),
      sizeof(TensorSliceProto_Extent),
      GOOGLE_PROTOBUF_GENERATED_MESSAGE_FIELD_OFFSET(TensorSliceProto_Extent, _internal_metadata_),
      GOOGLE_PROTOBUF_GENERATED_MESSAGE_FIELD_OFFSET(TensorSliceProto_Extent, _is_default_instance_));
}

namespace {

GOOGLE_PROTOBUF_DECLARE_ONCE(protobuf_AssignDescriptors_once_);
inline void protobuf_AssignDescriptorsOnce() {
  ::google::protobuf::GoogleOnceInit(&protobuf_AssignDescriptors_once_,
                 &protobuf_AssignDesc_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto);
}

void protobuf_RegisterTypes(const ::std::string&) {
  protobuf_AssignDescriptorsOnce();
  ::google::protobuf::MessageFactory::InternalRegisterGeneratedMessage(
      TensorSliceProto_descriptor_, &TensorSliceProto::default_instance());
  ::google::protobuf::MessageFactory::InternalRegisterGeneratedMessage(
      TensorSliceProto_Extent_descriptor_, &TensorSliceProto_Extent::default_instance());
}

}  // namespace

void protobuf_ShutdownFile_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto() {
  delete TensorSliceProto::default_instance_;
  delete TensorSliceProto_reflection_;
  delete TensorSliceProto_Extent::default_instance_;
  delete TensorSliceProto_Extent_default_oneof_instance_;
  delete TensorSliceProto_Extent_reflection_;
}

void protobuf_AddDesc_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto() {
  static bool already_here = false;
  if (already_here) return;
  already_here = true;
  GOOGLE_PROTOBUF_VERIFY_VERSION;

  ::google::protobuf::DescriptorPool::InternalAddGeneratedFile(
    "\n,tensorflow/core/framework/tensor_slice"
    ".proto\022\ntensorflow\"\200\001\n\020TensorSliceProto\022"
    "3\n\006extent\030\001 \003(\0132#.tensorflow.TensorSlice"
    "Proto.Extent\0327\n\006Extent\022\r\n\005start\030\001 \001(\003\022\020\n"
    "\006length\030\002 \001(\003H\000B\014\n\nhas_lengthB2\n\030org.ten"
    "sorflow.frameworkB\021TensorSliceProtosP\001\370\001"
    "\001b\006proto3", 249);
  ::google::protobuf::MessageFactory::InternalRegisterGeneratedFile(
    "tensorflow/core/framework/tensor_slice.proto", &protobuf_RegisterTypes);
  TensorSliceProto::default_instance_ = new TensorSliceProto();
  TensorSliceProto_Extent::default_instance_ = new TensorSliceProto_Extent();
  TensorSliceProto_Extent_default_oneof_instance_ = new TensorSliceProto_ExtentOneofInstance();
  TensorSliceProto::default_instance_->InitAsDefaultInstance();
  TensorSliceProto_Extent::default_instance_->InitAsDefaultInstance();
  ::google::protobuf::internal::OnShutdown(&protobuf_ShutdownFile_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto);
}

// Force AddDescriptors() to be called at static initialization time.
struct StaticDescriptorInitializer_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto {
  StaticDescriptorInitializer_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto() {
    protobuf_AddDesc_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto();
  }
} static_descriptor_initializer_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto_;

namespace {

static void MergeFromFail(int line) GOOGLE_ATTRIBUTE_COLD;
static void MergeFromFail(int line) {
  GOOGLE_CHECK(false) << __FILE__ << ":" << line;
}

}  // namespace

// ===================================================================

#if !defined(_MSC_VER) || _MSC_VER >= 1900
const int TensorSliceProto_Extent::kStartFieldNumber;
const int TensorSliceProto_Extent::kLengthFieldNumber;
#endif  // !defined(_MSC_VER) || _MSC_VER >= 1900

TensorSliceProto_Extent::TensorSliceProto_Extent()
  : ::google::protobuf::Message(), _internal_metadata_(NULL) {
  SharedCtor();
  // @@protoc_insertion_point(constructor:tensorflow.TensorSliceProto.Extent)
}

TensorSliceProto_Extent::TensorSliceProto_Extent(::google::protobuf::Arena* arena)
  : ::google::protobuf::Message(),
  _internal_metadata_(arena) {
  SharedCtor();
  RegisterArenaDtor(arena);
  // @@protoc_insertion_point(arena_constructor:tensorflow.TensorSliceProto.Extent)
}

void TensorSliceProto_Extent::InitAsDefaultInstance() {
  _is_default_instance_ = true;
  TensorSliceProto_Extent_default_oneof_instance_->length_ = GOOGLE_LONGLONG(0);
}

TensorSliceProto_Extent::TensorSliceProto_Extent(const TensorSliceProto_Extent& from)
  : ::google::protobuf::Message(),
    _internal_metadata_(NULL) {
  SharedCtor();
  MergeFrom(from);
  // @@protoc_insertion_point(copy_constructor:tensorflow.TensorSliceProto.Extent)
}

void TensorSliceProto_Extent::SharedCtor() {
    _is_default_instance_ = false;
  _cached_size_ = 0;
  start_ = GOOGLE_LONGLONG(0);
  clear_has_has_length();
}

TensorSliceProto_Extent::~TensorSliceProto_Extent() {
  // @@protoc_insertion_point(destructor:tensorflow.TensorSliceProto.Extent)
  SharedDtor();
}

void TensorSliceProto_Extent::SharedDtor() {
  if (GetArenaNoVirtual() != NULL) {
    return;
  }

  if (has_has_length()) {
    clear_has_length();
  }
  if (this != default_instance_) {
  }
}

void TensorSliceProto_Extent::ArenaDtor(void* object) {
  TensorSliceProto_Extent* _this = reinterpret_cast< TensorSliceProto_Extent* >(object);
  (void)_this;
}
void TensorSliceProto_Extent::RegisterArenaDtor(::google::protobuf::Arena* arena) {
}
void TensorSliceProto_Extent::SetCachedSize(int size) const {
  GOOGLE_SAFE_CONCURRENT_WRITES_BEGIN();
  _cached_size_ = size;
  GOOGLE_SAFE_CONCURRENT_WRITES_END();
}
const ::google::protobuf::Descriptor* TensorSliceProto_Extent::descriptor() {
  protobuf_AssignDescriptorsOnce();
  return TensorSliceProto_Extent_descriptor_;
}

const TensorSliceProto_Extent& TensorSliceProto_Extent::default_instance() {
  if (default_instance_ == NULL) protobuf_AddDesc_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto();
  return *default_instance_;
}

TensorSliceProto_Extent* TensorSliceProto_Extent::default_instance_ = NULL;

TensorSliceProto_Extent* TensorSliceProto_Extent::New(::google::protobuf::Arena* arena) const {
  return ::google::protobuf::Arena::CreateMessage<TensorSliceProto_Extent>(arena);
}

void TensorSliceProto_Extent::clear_has_length() {
// @@protoc_insertion_point(one_of_clear_start:tensorflow.TensorSliceProto.Extent)
  switch(has_length_case()) {
    case kLength: {
      // No need to clear
      break;
    }
    case HAS_LENGTH_NOT_SET: {
      break;
    }
  }
  _oneof_case_[0] = HAS_LENGTH_NOT_SET;
}

void TensorSliceProto_Extent::Clear() {
// @@protoc_insertion_point(message_clear_start:tensorflow.TensorSliceProto.Extent)
  start_ = GOOGLE_LONGLONG(0);
  clear_has_length();
}

bool TensorSliceProto_Extent::MergePartialFromCodedStream(
    ::google::protobuf::io::CodedInputStream* input) {
#define DO_(EXPRESSION) if (!GOOGLE_PREDICT_TRUE(EXPRESSION)) goto failure
  ::google::protobuf::uint32 tag;
  // @@protoc_insertion_point(parse_start:tensorflow.TensorSliceProto.Extent)
  for (;;) {
    ::std::pair< ::google::protobuf::uint32, bool> p = input->ReadTagWithCutoff(127);
    tag = p.first;
    if (!p.second) goto handle_unusual;
    switch (::google::protobuf::internal::WireFormatLite::GetTagFieldNumber(tag)) {
      // optional int64 start = 1;
      case 1: {
        if (tag == 8) {
          DO_((::google::protobuf::internal::WireFormatLite::ReadPrimitive<
                   ::google::protobuf::int64, ::google::protobuf::internal::WireFormatLite::TYPE_INT64>(
                 input, &start_)));

        } else {
          goto handle_unusual;
        }
        if (input->ExpectTag(16)) goto parse_length;
        break;
      }

      // optional int64 length = 2;
      case 2: {
        if (tag == 16) {
         parse_length:
          clear_has_length();
          DO_((::google::protobuf::internal::WireFormatLite::ReadPrimitive<
                   ::google::protobuf::int64, ::google::protobuf::internal::WireFormatLite::TYPE_INT64>(
                 input, &has_length_.length_)));
          set_has_length();
        } else {
          goto handle_unusual;
        }
        if (input->ExpectAtEnd()) goto success;
        break;
      }

      default: {
      handle_unusual:
        if (tag == 0 ||
            ::google::protobuf::internal::WireFormatLite::GetTagWireType(tag) ==
            ::google::protobuf::internal::WireFormatLite::WIRETYPE_END_GROUP) {
          goto success;
        }
        DO_(::google::protobuf::internal::WireFormatLite::SkipField(input, tag));
        break;
      }
    }
  }
success:
  // @@protoc_insertion_point(parse_success:tensorflow.TensorSliceProto.Extent)
  return true;
failure:
  // @@protoc_insertion_point(parse_failure:tensorflow.TensorSliceProto.Extent)
  return false;
#undef DO_
}

void TensorSliceProto_Extent::SerializeWithCachedSizes(
    ::google::protobuf::io::CodedOutputStream* output) const {
  // @@protoc_insertion_point(serialize_start:tensorflow.TensorSliceProto.Extent)
  // optional int64 start = 1;
  if (this->start() != 0) {
    ::google::protobuf::internal::WireFormatLite::WriteInt64(1, this->start(), output);
  }

  // optional int64 length = 2;
  if (has_length()) {
    ::google::protobuf::internal::WireFormatLite::WriteInt64(2, this->length(), output);
  }

  // @@protoc_insertion_point(serialize_end:tensorflow.TensorSliceProto.Extent)
}

::google::protobuf::uint8* TensorSliceProto_Extent::SerializeWithCachedSizesToArray(
    ::google::protobuf::uint8* target) const {
  // @@protoc_insertion_point(serialize_to_array_start:tensorflow.TensorSliceProto.Extent)
  // optional int64 start = 1;
  if (this->start() != 0) {
    target = ::google::protobuf::internal::WireFormatLite::WriteInt64ToArray(1, this->start(), target);
  }

  // optional int64 length = 2;
  if (has_length()) {
    target = ::google::protobuf::internal::WireFormatLite::WriteInt64ToArray(2, this->length(), target);
  }

  // @@protoc_insertion_point(serialize_to_array_end:tensorflow.TensorSliceProto.Extent)
  return target;
}

int TensorSliceProto_Extent::ByteSize() const {
// @@protoc_insertion_point(message_byte_size_start:tensorflow.TensorSliceProto.Extent)
  int total_size = 0;

  // optional int64 start = 1;
  if (this->start() != 0) {
    total_size += 1 +
      ::google::protobuf::internal::WireFormatLite::Int64Size(
        this->start());
  }

  switch (has_length_case()) {
    // optional int64 length = 2;
    case kLength: {
      total_size += 1 +
        ::google::protobuf::internal::WireFormatLite::Int64Size(
          this->length());
      break;
    }
    case HAS_LENGTH_NOT_SET: {
      break;
    }
  }
  GOOGLE_SAFE_CONCURRENT_WRITES_BEGIN();
  _cached_size_ = total_size;
  GOOGLE_SAFE_CONCURRENT_WRITES_END();
  return total_size;
}

void TensorSliceProto_Extent::MergeFrom(const ::google::protobuf::Message& from) {
// @@protoc_insertion_point(generalized_merge_from_start:tensorflow.TensorSliceProto.Extent)
  if (GOOGLE_PREDICT_FALSE(&from == this)) MergeFromFail(__LINE__);
  const TensorSliceProto_Extent* source =
      ::google::protobuf::internal::DynamicCastToGenerated<const TensorSliceProto_Extent>(
          &from);
  if (source == NULL) {
  // @@protoc_insertion_point(generalized_merge_from_cast_fail:tensorflow.TensorSliceProto.Extent)
    ::google::protobuf::internal::ReflectionOps::Merge(from, this);
  } else {
  // @@protoc_insertion_point(generalized_merge_from_cast_success:tensorflow.TensorSliceProto.Extent)
    MergeFrom(*source);
  }
}

void TensorSliceProto_Extent::MergeFrom(const TensorSliceProto_Extent& from) {
// @@protoc_insertion_point(class_specific_merge_from_start:tensorflow.TensorSliceProto.Extent)
  if (GOOGLE_PREDICT_FALSE(&from == this)) MergeFromFail(__LINE__);
  switch (from.has_length_case()) {
    case kLength: {
      set_length(from.length());
      break;
    }
    case HAS_LENGTH_NOT_SET: {
      break;
    }
  }
  if (from.start() != 0) {
    set_start(from.start());
  }
}

void TensorSliceProto_Extent::CopyFrom(const ::google::protobuf::Message& from) {
// @@protoc_insertion_point(generalized_copy_from_start:tensorflow.TensorSliceProto.Extent)
  if (&from == this) return;
  Clear();
  MergeFrom(from);
}

void TensorSliceProto_Extent::CopyFrom(const TensorSliceProto_Extent& from) {
// @@protoc_insertion_point(class_specific_copy_from_start:tensorflow.TensorSliceProto.Extent)
  if (&from == this) return;
  Clear();
  MergeFrom(from);
}

bool TensorSliceProto_Extent::IsInitialized() const {

  return true;
}

void TensorSliceProto_Extent::Swap(TensorSliceProto_Extent* other) {
  if (other == this) return;
  if (GetArenaNoVirtual() == other->GetArenaNoVirtual()) {
    InternalSwap(other);
  } else {
    TensorSliceProto_Extent temp;
    temp.MergeFrom(*this);
    CopyFrom(*other);
    other->CopyFrom(temp);
  }
}
void TensorSliceProto_Extent::UnsafeArenaSwap(TensorSliceProto_Extent* other) {
  if (other == this) return;
  GOOGLE_DCHECK(GetArenaNoVirtual() == other->GetArenaNoVirtual());
  InternalSwap(other);
}
void TensorSliceProto_Extent::InternalSwap(TensorSliceProto_Extent* other) {
  std::swap(start_, other->start_);
  std::swap(has_length_, other->has_length_);
  std::swap(_oneof_case_[0], other->_oneof_case_[0]);
  _internal_metadata_.Swap(&other->_internal_metadata_);
  std::swap(_cached_size_, other->_cached_size_);
}

::google::protobuf::Metadata TensorSliceProto_Extent::GetMetadata() const {
  protobuf_AssignDescriptorsOnce();
  ::google::protobuf::Metadata metadata;
  metadata.descriptor = TensorSliceProto_Extent_descriptor_;
  metadata.reflection = TensorSliceProto_Extent_reflection_;
  return metadata;
}

// -------------------------------------------------------------------

#if !defined(_MSC_VER) || _MSC_VER >= 1900
const int TensorSliceProto::kExtentFieldNumber;
#endif  // !defined(_MSC_VER) || _MSC_VER >= 1900

TensorSliceProto::TensorSliceProto()
  : ::google::protobuf::Message(), _internal_metadata_(NULL) {
  SharedCtor();
  // @@protoc_insertion_point(constructor:tensorflow.TensorSliceProto)
}

TensorSliceProto::TensorSliceProto(::google::protobuf::Arena* arena)
  : ::google::protobuf::Message(),
  _internal_metadata_(arena),
  extent_(arena) {
  SharedCtor();
  RegisterArenaDtor(arena);
  // @@protoc_insertion_point(arena_constructor:tensorflow.TensorSliceProto)
}

void TensorSliceProto::InitAsDefaultInstance() {
  _is_default_instance_ = true;
}

TensorSliceProto::TensorSliceProto(const TensorSliceProto& from)
  : ::google::protobuf::Message(),
    _internal_metadata_(NULL) {
  SharedCtor();
  MergeFrom(from);
  // @@protoc_insertion_point(copy_constructor:tensorflow.TensorSliceProto)
}

void TensorSliceProto::SharedCtor() {
    _is_default_instance_ = false;
  _cached_size_ = 0;
}

TensorSliceProto::~TensorSliceProto() {
  // @@protoc_insertion_point(destructor:tensorflow.TensorSliceProto)
  SharedDtor();
}

void TensorSliceProto::SharedDtor() {
  if (GetArenaNoVirtual() != NULL) {
    return;
  }

  if (this != default_instance_) {
  }
}

void TensorSliceProto::ArenaDtor(void* object) {
  TensorSliceProto* _this = reinterpret_cast< TensorSliceProto* >(object);
  (void)_this;
}
void TensorSliceProto::RegisterArenaDtor(::google::protobuf::Arena* arena) {
}
void TensorSliceProto::SetCachedSize(int size) const {
  GOOGLE_SAFE_CONCURRENT_WRITES_BEGIN();
  _cached_size_ = size;
  GOOGLE_SAFE_CONCURRENT_WRITES_END();
}
const ::google::protobuf::Descriptor* TensorSliceProto::descriptor() {
  protobuf_AssignDescriptorsOnce();
  return TensorSliceProto_descriptor_;
}

const TensorSliceProto& TensorSliceProto::default_instance() {
  if (default_instance_ == NULL) protobuf_AddDesc_tensorflow_2fcore_2fframework_2ftensor_5fslice_2eproto();
  return *default_instance_;
}

TensorSliceProto* TensorSliceProto::default_instance_ = NULL;

TensorSliceProto* TensorSliceProto::New(::google::protobuf::Arena* arena) const {
  return ::google::protobuf::Arena::CreateMessage<TensorSliceProto>(arena);
}

void TensorSliceProto::Clear() {
// @@protoc_insertion_point(message_clear_start:tensorflow.TensorSliceProto)
  extent_.Clear();
}

bool TensorSliceProto::MergePartialFromCodedStream(
    ::google::protobuf::io::CodedInputStream* input) {
#define DO_(EXPRESSION) if (!GOOGLE_PREDICT_TRUE(EXPRESSION)) goto failure
  ::google::protobuf::uint32 tag;
  // @@protoc_insertion_point(parse_start:tensorflow.TensorSliceProto)
  for (;;) {
    ::std::pair< ::google::protobuf::uint32, bool> p = input->ReadTagWithCutoff(127);
    tag = p.first;
    if (!p.second) goto handle_unusual;
    switch (::google::protobuf::internal::WireFormatLite::GetTagFieldNumber(tag)) {
      // repeated .tensorflow.TensorSliceProto.Extent extent = 1;
      case 1: {
        if (tag == 10) {
          DO_(input->IncrementRecursionDepth());
         parse_loop_extent:
          DO_(::google::protobuf::internal::WireFormatLite::ReadMessageNoVirtualNoRecursionDepth(
                input, add_extent()));
        } else {
          goto handle_unusual;
        }
        if (input->ExpectTag(10)) goto parse_loop_extent;
        input->UnsafeDecrementRecursionDepth();
        if (input->ExpectAtEnd()) goto success;
        break;
      }

      default: {
      handle_unusual:
        if (tag == 0 ||
            ::google::protobuf::internal::WireFormatLite::GetTagWireType(tag) ==
            ::google::protobuf::internal::WireFormatLite::WIRETYPE_END_GROUP) {
          goto success;
        }
        DO_(::google::protobuf::internal::WireFormatLite::SkipField(input, tag));
        break;
      }
    }
  }
success:
  // @@protoc_insertion_point(parse_success:tensorflow.TensorSliceProto)
  return true;
failure:
  // @@protoc_insertion_point(parse_failure:tensorflow.TensorSliceProto)
  return false;
#undef DO_
}

void TensorSliceProto::SerializeWithCachedSizes(
    ::google::protobuf::io::CodedOutputStream* output) const {
  // @@protoc_insertion_point(serialize_start:tensorflow.TensorSliceProto)
  // repeated .tensorflow.TensorSliceProto.Extent extent = 1;
  for (unsigned int i = 0, n = this->extent_size(); i < n; i++) {
    ::google::protobuf::internal::WireFormatLite::WriteMessageMaybeToArray(
      1, this->extent(i), output);
  }

  // @@protoc_insertion_point(serialize_end:tensorflow.TensorSliceProto)
}

::google::protobuf::uint8* TensorSliceProto::SerializeWithCachedSizesToArray(
    ::google::protobuf::uint8* target) const {
  // @@protoc_insertion_point(serialize_to_array_start:tensorflow.TensorSliceProto)
  // repeated .tensorflow.TensorSliceProto.Extent extent = 1;
  for (unsigned int i = 0, n = this->extent_size(); i < n; i++) {
    target = ::google::protobuf::internal::WireFormatLite::
      WriteMessageNoVirtualToArray(
        1, this->extent(i), target);
  }

  // @@protoc_insertion_point(serialize_to_array_end:tensorflow.TensorSliceProto)
  return target;
}

int TensorSliceProto::ByteSize() const {
// @@protoc_insertion_point(message_byte_size_start:tensorflow.TensorSliceProto)
  int total_size = 0;

  // repeated .tensorflow.TensorSliceProto.Extent extent = 1;
  total_size += 1 * this->extent_size();
  for (int i = 0; i < this->extent_size(); i++) {
    total_size +=
      ::google::protobuf::internal::WireFormatLite::MessageSizeNoVirtual(
        this->extent(i));
  }

  GOOGLE_SAFE_CONCURRENT_WRITES_BEGIN();
  _cached_size_ = total_size;
  GOOGLE_SAFE_CONCURRENT_WRITES_END();
  return total_size;
}

void TensorSliceProto::MergeFrom(const ::google::protobuf::Message& from) {
// @@protoc_insertion_point(generalized_merge_from_start:tensorflow.TensorSliceProto)
  if (GOOGLE_PREDICT_FALSE(&from == this)) MergeFromFail(__LINE__);
  const TensorSliceProto* source =
      ::google::protobuf::internal::DynamicCastToGenerated<const TensorSliceProto>(
          &from);
  if (source == NULL) {
  // @@protoc_insertion_point(generalized_merge_from_cast_fail:tensorflow.TensorSliceProto)
    ::google::protobuf::internal::ReflectionOps::Merge(from, this);
  } else {
  // @@protoc_insertion_point(generalized_merge_from_cast_success:tensorflow.TensorSliceProto)
    MergeFrom(*source);
  }
}

void TensorSliceProto::MergeFrom(const TensorSliceProto& from) {
// @@protoc_insertion_point(class_specific_merge_from_start:tensorflow.TensorSliceProto)
  if (GOOGLE_PREDICT_FALSE(&from == this)) MergeFromFail(__LINE__);
  extent_.MergeFrom(from.extent_);
}

void TensorSliceProto::CopyFrom(const ::google::protobuf::Message& from) {
// @@protoc_insertion_point(generalized_copy_from_start:tensorflow.TensorSliceProto)
  if (&from == this) return;
  Clear();
  MergeFrom(from);
}

void TensorSliceProto::CopyFrom(const TensorSliceProto& from) {
// @@protoc_insertion_point(class_specific_copy_from_start:tensorflow.TensorSliceProto)
  if (&from == this) return;
  Clear();
  MergeFrom(from);
}

bool TensorSliceProto::IsInitialized() const {

  return true;
}

void TensorSliceProto::Swap(TensorSliceProto* other) {
  if (other == this) return;
  if (GetArenaNoVirtual() == other->GetArenaNoVirtual()) {
    InternalSwap(other);
  } else {
    TensorSliceProto temp;
    temp.MergeFrom(*this);
    CopyFrom(*other);
    other->CopyFrom(temp);
  }
}
void TensorSliceProto::UnsafeArenaSwap(TensorSliceProto* other) {
  if (other == this) return;
  GOOGLE_DCHECK(GetArenaNoVirtual() == other->GetArenaNoVirtual());
  InternalSwap(other);
}
void TensorSliceProto::InternalSwap(TensorSliceProto* other) {
  extent_.UnsafeArenaSwap(&other->extent_);
  _internal_metadata_.Swap(&other->_internal_metadata_);
  std::swap(_cached_size_, other->_cached_size_);
}

::google::protobuf::Metadata TensorSliceProto::GetMetadata() const {
  protobuf_AssignDescriptorsOnce();
  ::google::protobuf::Metadata metadata;
  metadata.descriptor = TensorSliceProto_descriptor_;
  metadata.reflection = TensorSliceProto_reflection_;
  return metadata;
}

#if PROTOBUF_INLINE_NOT_IN_HEADERS
// TensorSliceProto_Extent

// optional int64 start = 1;
void TensorSliceProto_Extent::clear_start() {
  start_ = GOOGLE_LONGLONG(0);
}
 ::google::protobuf::int64 TensorSliceProto_Extent::start() const {
  // @@protoc_insertion_point(field_get:tensorflow.TensorSliceProto.Extent.start)
  return start_;
}
 void TensorSliceProto_Extent::set_start(::google::protobuf::int64 value) {

  start_ = value;
  // @@protoc_insertion_point(field_set:tensorflow.TensorSliceProto.Extent.start)
}

// optional int64 length = 2;
bool TensorSliceProto_Extent::has_length() const {
  return has_length_case() == kLength;
}
void TensorSliceProto_Extent::set_has_length() {
  _oneof_case_[0] = kLength;
}
void TensorSliceProto_Extent::clear_length() {
  if (has_length()) {
    has_length_.length_ = GOOGLE_LONGLONG(0);
    clear_has_has_length();
  }
}
 ::google::protobuf::int64 TensorSliceProto_Extent::length() const {
  // @@protoc_insertion_point(field_get:tensorflow.TensorSliceProto.Extent.length)
  if (has_length()) {
    return has_length_.length_;
  }
  return GOOGLE_LONGLONG(0);
}
 void TensorSliceProto_Extent::set_length(::google::protobuf::int64 value) {
  if (!has_length()) {
    clear_has_length();
    set_has_length();
  }
  has_length_.length_ = value;
  // @@protoc_insertion_point(field_set:tensorflow.TensorSliceProto.Extent.length)
}

bool TensorSliceProto_Extent::has_has_length() const {
  return has_length_case() != HAS_LENGTH_NOT_SET;
}
void TensorSliceProto_Extent::clear_has_has_length() {
  _oneof_case_[0] = HAS_LENGTH_NOT_SET;
}
TensorSliceProto_Extent::HasLengthCase TensorSliceProto_Extent::has_length_case() const {
  return TensorSliceProto_Extent::HasLengthCase(_oneof_case_[0]);
}
// -------------------------------------------------------------------

// TensorSliceProto

// repeated .tensorflow.TensorSliceProto.Extent extent = 1;
int TensorSliceProto::extent_size() const {
  return extent_.size();
}
void TensorSliceProto::clear_extent() {
  extent_.Clear();
}
const ::tensorflow::TensorSliceProto_Extent& TensorSliceProto::extent(int index) const {
  // @@protoc_insertion_point(field_get:tensorflow.TensorSliceProto.extent)
  return extent_.Get(index);
}
::tensorflow::TensorSliceProto_Extent* TensorSliceProto::mutable_extent(int index) {
  // @@protoc_insertion_point(field_mutable:tensorflow.TensorSliceProto.extent)
  return extent_.Mutable(index);
}
::tensorflow::TensorSliceProto_Extent* TensorSliceProto::add_extent() {
  // @@protoc_insertion_point(field_add:tensorflow.TensorSliceProto.extent)
  return extent_.Add();
}
::google::protobuf::RepeatedPtrField< ::tensorflow::TensorSliceProto_Extent >*
TensorSliceProto::mutable_extent() {
  // @@protoc_insertion_point(field_mutable_list:tensorflow.TensorSliceProto.extent)
  return &extent_;
}
const ::google::protobuf::RepeatedPtrField< ::tensorflow::TensorSliceProto_Extent >&
TensorSliceProto::extent() const {
  // @@protoc_insertion_point(field_list:tensorflow.TensorSliceProto.extent)
  return extent_;
}

#endif  // PROTOBUF_INLINE_NOT_IN_HEADERS

// @@protoc_insertion_point(namespace_scope)

}  // namespace tensorflow

// @@protoc_insertion_point(global_scope)

redpajama_set_name: RedPajamaGithub


text:
"""Tests for mux module."""

import Queue
import logging
import optparse
import unittest
import struct
import sys

import set_sys_path  # Update sys.path to locate mod_pywebsocket module.

from mod_pywebsocket import common
from mod_pywebsocket import mux
from mod_pywebsocket._stream_base import ConnectionTerminatedException
from mod_pywebsocket._stream_hybi import Stream
from mod_pywebsocket._stream_hybi import StreamOptions
from mod_pywebsocket._stream_hybi import create_binary_frame
from mod_pywebsocket._stream_hybi import parse_frame

import mock

class _OutgoingChannelData(object):
    def __init__(self):
        self.messages = []
        self.control_messages = []

        self.current_opcode = None
        self.pending_fragments = []

class _MockMuxConnection(mock.MockBlockingConn):
    """Mock class of mod_python connection for mux."""

    def __init__(self):
        mock.MockBlockingConn.__init__(self)
        self._control_blocks = []
        self._channel_data = {}

        self._current_opcode = None
        self._pending_fragments = []

    def write(self, data):
        """Override MockBlockingConn.write."""

        self._current_data = data
        self._position = 0

        def _receive_bytes(length):
            if self._position + length > len(self._current_data):
                raise ConnectionTerminatedException(
                    'Failed to receive %d bytes from encapsulated '
                    'frame' % length)
            data = self._current_data[self._position:self._position+length]
            self._position += length
            return data

        opcode, payload, fin, rsv1, rsv2, rsv3 = (
            parse_frame(_receive_bytes, unmask_receive=False))

        self._pending_fragments.append(payload)

        if self._current_opcode is None:
            if opcode == common.OPCODE_CONTINUATION:
                raise Exception('Sending invalid continuation opcode')
            self._current_opcode = opcode
        else:
            if opcode != common.OPCODE_CONTINUATION:
                raise Exception('Sending invalid opcode %d' % opcode)
        if not fin:
            return

        inner_frame_data = ''.join(self._pending_fragments)
        self._pending_fragments = []
        self._current_opcode = None

        parser = mux._MuxFramePayloadParser(inner_frame_data)
        channel_id = parser.read_channel_id()
        if channel_id == mux._CONTROL_CHANNEL_ID:
            self._control_blocks.append(parser.remaining_data())
            return

        if not channel_id in self._channel_data:
            self._channel_data[channel_id] = _OutgoingChannelData()
        channel_data = self._channel_data[channel_id]

        (inner_fin, inner_rsv1, inner_rsv2, inner_rsv3, inner_opcode,
         inner_payload) = parser.read_inner_frame()
        channel_data.pending_fragments.append(inner_payload)

        if channel_data.current_opcode is None:
            if inner_opcode == common.OPCODE_CONTINUATION:
                raise Exception('Sending invalid continuation opcode')
            channel_data.current_opcode = inner_opcode
        else:
            if inner_opcode != common.OPCODE_CONTINUATION:
                raise Exception('Sending invalid opcode %d' % inner_opcode)
        if not inner_fin:
            return

        message = ''.join(channel_data.pending_fragments)
        channel_data.pending_fragments = []

        if (channel_data.current_opcode == common.OPCODE_TEXT or
            channel_data.current_opcode == common.OPCODE_BINARY):
            channel_data.messages.append(message)
        else:
            channel_data.control_messages.append(
                {'opcode': channel_data.current_opcode,
                 'message': message})
        channel_data.current_opcode = None

    def get_written_control_blocks(self):
        return self._control_blocks

    def get_written_messages(self, channel_id):
        return self._channel_data[channel_id].messages

    def get_written_control_messages(self, channel_id):
        return self._channel_data[channel_id].control_messages

class _ChannelEvent(object):
    """A structure that records channel events."""

    def __init__(self):
        self.messages = []
        self.exception = None
        self.client_initiated_closing = False

class _MuxMockDispatcher(object):
    """Mock class of dispatch.Dispatcher for mux."""

    def __init__(self):
        self.channel_events = {}

    def do_extra_handshake(self, request):
        pass

    def _do_echo(self, request, channel_events):
        while True:
            message = request.ws_stream.receive_message()
            if message == None:
                channel_events.client_initiated_closing = True
                return
            if message == 'Goodbye':
                return
            channel_events.messages.append(message)
            # echo back
            request.ws_stream.send_message(message)

    def _do_ping(self, request, channel_events):
        request.ws_stream.send_ping('Ping!')

    def transfer_data(self, request):
        self.channel_events[request.channel_id] = _ChannelEvent()

        try:
            # Note: more handler will be added.
            if request.uri.endswith('echo'):
                self._do_echo(request,
                              self.channel_events[request.channel_id])
            elif request.uri.endswith('ping'):
                self._do_ping(request,
                              self.channel_events[request.channel_id])
            else:
                raise ValueError('Cannot handle path %r' % request.path)
            if not request.server_terminated:
                request.ws_stream.close_connection()
        except ConnectionTerminatedException, e:
            self.channel_events[request.channel_id].exception = e
        except Exception, e:
            self.channel_events[request.channel_id].exception = e
            raise

def _create_mock_request():
    headers = {'Host': 'server.example.com',
               'Upgrade': 'websocket',
               'Connection': 'Upgrade',
               'Sec-WebSocket-Key': 'dGhlIHNhbXBsZSBub25jZQ==',
               'Sec-WebSocket-Version': '13',
               'Origin': 'http://example.com'}
    request = mock.MockRequest(uri='/echo',
                               headers_in=headers,
                               connection=_MockMuxConnection())
    request.ws_stream = Stream(request, options=StreamOptions())
    request.mux = True
    request.mux_extensions = []
    request.mux_quota = 8 * 1024
    return request

def _create_add_channel_request_frame(channel_id, encoding, encoded_handshake):
    if encoding != 0 and encoding != 1:
        raise ValueError('Invalid encoding')
    block = mux._create_control_block_length_value(
               channel_id, mux._MUX_OPCODE_ADD_CHANNEL_REQUEST, encoding,
               encoded_handshake)
    payload = mux._encode_channel_id(mux._CONTROL_CHANNEL_ID) + block
    return create_binary_frame(payload, mask=True)

def _create_logical_frame(channel_id, message, opcode=common.OPCODE_BINARY,
                          mask=True):
    bits = chr(0x80 | opcode)
    payload = mux._encode_channel_id(channel_id) + bits + message
    return create_binary_frame(payload, mask=mask)

def _create_request_header(path='/echo'):
    return (
        'GET %s HTTP/1.1\r\n'
        'Host: server.example.com\r\n'
        'Upgrade: websocket\r\n'
        'Connection: Upgrade\r\n'
        'Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n'
        'Sec-WebSocket-Version: 13\r\n'
        'Origin: http://example.com\r\n'
        '\r\n') % path

class MuxTest(unittest.TestCase):
    """A unittest for mux module."""

    def test_channel_id_decode(self):
        data = '\x00\x01\xbf\xff\xdf\xff\xff\xff\xff\xff\xff'
        parser = mux._MuxFramePayloadParser(data)
        channel_id = parser.read_channel_id()
        self.assertEqual(0, channel_id)
        channel_id = parser.read_channel_id()
        self.assertEqual(1, channel_id)
        channel_id = parser.read_channel_id()
        self.assertEqual(2 ** 14 - 1, channel_id)
        channel_id = parser.read_channel_id()
        self.assertEqual(2 ** 21 - 1, channel_id)
        channel_id = parser.read_channel_id()
        self.assertEqual(2 ** 29 - 1, channel_id)
        self.assertEqual(len(data), parser._read_position)

    def test_channel_id_encode(self):
        encoded = mux._encode_channel_id(0)
        self.assertEqual('\x00', encoded)
        encoded = mux._encode_channel_id(2 ** 14 - 1)
        self.assertEqual('\xbf\xff', encoded)
        encoded = mux._encode_channel_id(2 ** 14)
        self.assertEqual('\xc0@\x00', encoded)
        encoded = mux._encode_channel_id(2 ** 21 - 1)
        self.assertEqual('\xdf\xff\xff', encoded)
        encoded = mux._encode_channel_id(2 ** 21)
        self.assertEqual('\xe0 \x00\x00', encoded)
        encoded = mux._encode_channel_id(2 ** 29 - 1)
        self.assertEqual('\xff\xff\xff\xff', encoded)
        # channel_id is too large
        self.assertRaises(ValueError,
                          mux._encode_channel_id,
                          2 ** 29)

    def test_create_control_block_length_value(self):
        data = 'Hello, world!'
        block = mux._create_control_block_length_value(
            channel_id=1, opcode=mux._MUX_OPCODE_ADD_CHANNEL_REQUEST,
            flags=0x7, value=data)
        expected = '\x1c\x01\x0dHello, world!'
        self.assertEqual(expected, block)

        data = 'a' * (2 ** 8)
        block = mux._create_control_block_length_value(
            channel_id=2, opcode=mux._MUX_OPCODE_ADD_CHANNEL_RESPONSE,
            flags=0x0, value=data)
        expected = '\x21\x02\x01\x00' + data
        self.assertEqual(expected, block)

        data = 'b' * (2 ** 16)
        block = mux._create_control_block_length_value(
            channel_id=3, opcode=mux._MUX_OPCODE_DROP_CHANNEL,
            flags=0x0, value=data)
        expected = '\x62\x03\x01\x00\x00' + data
        self.assertEqual(expected, block)

    def test_read_control_blocks(self):
        data = ('\x00\x01\00'
                '\x61\x02\x01\x00%s'
                '\x0a\x03\x01\x00\x00%s'
                '\x63\x04\x01\x00\x00\x00%s') % (
            'a' * 0x0100, 'b' * 0x010000, 'c' * 0x01000000)
        parser = mux._MuxFramePayloadParser(data)
        blocks = list(parser.read_control_blocks())
        self.assertEqual(4, len(blocks))

        self.assertEqual(mux._MUX_OPCODE_ADD_CHANNEL_REQUEST, blocks[0].opcode)
        self.assertEqual(0, blocks[0].encoding)
        self.assertEqual(0, len(blocks[0].encoded_handshake))

        self.assertEqual(mux._MUX_OPCODE_DROP_CHANNEL, blocks[1].opcode)
        self.assertEqual(0, blocks[1].mux_error)
        self.assertEqual(0x0100, len(blocks[1].reason))

        self.assertEqual(mux._MUX_OPCODE_ADD_CHANNEL_REQUEST, blocks[2].opcode)
        self.assertEqual(2, blocks[2].encoding)
        self.assertEqual(0x010000, len(blocks[2].encoded_handshake))

        self.assertEqual(mux._MUX_OPCODE_DROP_CHANNEL, blocks[3].opcode)
        self.assertEqual(0, blocks[3].mux_error)
        self.assertEqual(0x01000000, len(blocks[3].reason))

        self.assertEqual(len(data), parser._read_position)

    def test_create_add_channel_response(self):
        data = mux._create_add_channel_response(channel_id=1,
                                                encoded_handshake='FooBar',
                                                encoding=0,
                                                rejected=False)
        self.assertEqual('\x82\x0a\x00\x20\x01\x06FooBar', data)

        data = mux._create_add_channel_response(channel_id=2,
                                                encoded_handshake='Hello',
                                                encoding=1,
                                                rejected=True)
        self.assertEqual('\x82\x09\x00\x34\x02\x05Hello', data)

    def test_drop_channel(self):
        data = mux._create_drop_channel(channel_id=1,
                                        reason='',
                                        mux_error=False)
        self.assertEqual('\x82\x04\x00\x60\x01\x00', data)

        data = mux._create_drop_channel(channel_id=1,
                                        reason='error',
                                        mux_error=True)
        self.assertEqual('\x82\x09\x00\x70\x01\x05error', data)

        # reason must be empty if mux_error is False.
        self.assertRaises(ValueError,
                          mux._create_drop_channel,
                          1, 'FooBar', False)

    def test_parse_request_text(self):
        request_text = _create_request_header()
        command, path, version, headers = mux._parse_request_text(request_text)
        self.assertEqual('GET', command)
        self.assertEqual('/echo', path)
        self.assertEqual('HTTP/1.1', version)
        self.assertEqual(6, len(headers))
        self.assertEqual('server.example.com', headers['Host'])
        self.assertEqual('websocket', headers['Upgrade'])
        self.assertEqual('Upgrade', headers['Connection'])
        self.assertEqual('dGhlIHNhbXBsZSBub25jZQ==',
                         headers['Sec-WebSocket-Key'])
        self.assertEqual('13', headers['Sec-WebSocket-Version'])
        self.assertEqual('http://example.com', headers['Origin'])

class MuxHandlerTest(unittest.TestCase):

    def test_add_channel(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(mux._INITIAL_NUMBER_OF_CHANNEL_SLOTS,
                                      mux._INITIAL_QUOTA_FOR_CLIENT)

        encoded_handshake = _create_request_header(path='/echo')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        flow_control = mux._create_flow_control(channel_id=2,
                                                replenished_quota=5,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        encoded_handshake = _create_request_header(path='/echo')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=3, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        flow_control = mux._create_flow_control(channel_id=3,
                                                replenished_quota=5,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=2, message='Hello'))
        request.connection.put_bytes(
            _create_logical_frame(channel_id=3, message='World'))
        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))
        request.connection.put_bytes(
            _create_logical_frame(channel_id=2, message='Goodbye'))
        request.connection.put_bytes(
            _create_logical_frame(channel_id=3, message='Goodbye'))

        mux_handler.wait_until_done(timeout=2)

        self.assertEqual([], dispatcher.channel_events[1].messages)
        self.assertEqual(['Hello'], dispatcher.channel_events[2].messages)
        self.assertEqual(['World'], dispatcher.channel_events[3].messages)
        # Channel 2
        messages = request.connection.get_written_messages(2)
        self.assertEqual(1, len(messages))
        self.assertEqual('Hello', messages[0])
        # Channel 3
        messages = request.connection.get_written_messages(3)
        self.assertEqual(1, len(messages))
        self.assertEqual('World', messages[0])
        control_blocks = request.connection.get_written_control_blocks()
        # There should be 8 control blocks:
        #   - 1 NewChannelSlot
        #   - 2 AddChannelResponses for channel id 2 and 3
        #   - 6 FlowControls for channel id 1 (initialize), 'Hello', 'World',
        #     and 3 'Goodbye's
        self.assertEqual(9, len(control_blocks))

    def test_add_channel_incomplete_handshake(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(mux._INITIAL_NUMBER_OF_CHANNEL_SLOTS,
                                      mux._INITIAL_QUOTA_FOR_CLIENT)

        incomplete_encoded_handshake = 'GET /echo HTTP/1.1'
        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=incomplete_encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))

        mux_handler.wait_until_done(timeout=2)

        self.assertTrue(1 in dispatcher.channel_events)
        self.assertTrue(not 2 in dispatcher.channel_events)

    def test_add_channel_invalid_version_handshake(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(mux._INITIAL_NUMBER_OF_CHANNEL_SLOTS,
                                      mux._INITIAL_QUOTA_FOR_CLIENT)

        encoded_handshake = (
            'GET /echo HTTP/1.1\r\n'
            'Host: example.com\r\n'
            'Connection: Upgrade\r\n'
            'Sec-WebSocket-Key2: 12998 5 Y3 1  .P00\r\n'
            'Sec-WebSocket-Protocol: sample\r\n'
            'Upgrade: WebSocket\r\n'
            'Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5\r\n'
            'Origin: http://example.com\r\n'
            '\r\n'
            '^n:ds[4U')

        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))

        mux_handler.wait_until_done(timeout=2)

        self.assertTrue(1 in dispatcher.channel_events)
        self.assertTrue(not 2 in dispatcher.channel_events)

    def test_receive_drop_channel(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(mux._INITIAL_NUMBER_OF_CHANNEL_SLOTS,
                                      mux._INITIAL_QUOTA_FOR_CLIENT)

        encoded_handshake = _create_request_header(path='/echo')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        drop_channel = mux._create_drop_channel(channel_id=2,
                                                outer_frame_mask=True)
        request.connection.put_bytes(drop_channel)

        # Terminate implicitly opened channel.
        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))

        mux_handler.wait_until_done(timeout=2)

        exception = dispatcher.channel_events[2].exception
        self.assertTrue(exception.__class__ == ConnectionTerminatedException)

    def test_receive_ping_frame(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(mux._INITIAL_NUMBER_OF_CHANNEL_SLOTS,
                                      mux._INITIAL_QUOTA_FOR_CLIENT)

        encoded_handshake = _create_request_header(path='/echo')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        flow_control = mux._create_flow_control(channel_id=2,
                                                replenished_quota=12,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        ping_frame = _create_logical_frame(channel_id=2,
                                           message='Hello World!',
                                           opcode=common.OPCODE_PING)
        request.connection.put_bytes(ping_frame)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))
        request.connection.put_bytes(
            _create_logical_frame(channel_id=2, message='Goodbye'))

        mux_handler.wait_until_done(timeout=2)

        messages = request.connection.get_written_control_messages(2)
        self.assertEqual(common.OPCODE_PONG, messages[0]['opcode'])
        self.assertEqual('Hello World!', messages[0]['message'])

    def test_send_ping(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(mux._INITIAL_NUMBER_OF_CHANNEL_SLOTS,
                                      mux._INITIAL_QUOTA_FOR_CLIENT)

        encoded_handshake = _create_request_header(path='/ping')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        flow_control = mux._create_flow_control(channel_id=2,
                                                replenished_quota=5,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))

        mux_handler.wait_until_done(timeout=2)

        messages = request.connection.get_written_control_messages(2)
        self.assertEqual(common.OPCODE_PING, messages[0]['opcode'])
        self.assertEqual('Ping!', messages[0]['message'])

    def test_two_flow_control(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(mux._INITIAL_NUMBER_OF_CHANNEL_SLOTS,
                                      mux._INITIAL_QUOTA_FOR_CLIENT)

        encoded_handshake = _create_request_header(path='/echo')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        # Replenish 5 bytes.
        flow_control = mux._create_flow_control(channel_id=2,
                                                replenished_quota=5,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        # Send 10 bytes. The server will try echo back 10 bytes.
        request.connection.put_bytes(
            _create_logical_frame(channel_id=2, message='HelloWorld'))

        # Replenish 5 bytes again.
        flow_control = mux._create_flow_control(channel_id=2,
                                                replenished_quota=5,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))
        request.connection.put_bytes(
            _create_logical_frame(channel_id=2, message='Goodbye'))

        mux_handler.wait_until_done(timeout=2)

        messages = request.connection.get_written_messages(2)
        self.assertEqual(['HelloWorld'], messages)

    def test_no_send_quota_on_server(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(mux._INITIAL_NUMBER_OF_CHANNEL_SLOTS,
                                      mux._INITIAL_QUOTA_FOR_CLIENT)

        encoded_handshake = _create_request_header(path='/echo')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=2, message='HelloWorld'))

        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))

        mux_handler.wait_until_done(timeout=1)

        # No message should be sent on channel 2.
        self.assertRaises(KeyError,
                          request.connection.get_written_messages,
                          2)

    def test_quota_violation_by_client(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(mux._INITIAL_NUMBER_OF_CHANNEL_SLOTS, 0)

        encoded_handshake = _create_request_header(path='/echo')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=2, message='HelloWorld'))

        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))

        mux_handler.wait_until_done(timeout=2)
        control_blocks = request.connection.get_written_control_blocks()
        # The first block is FlowControl for channel id 1.
        # The next two blocks are NewChannelSlot and AddChannelResponse.
        # The 4th block or the last block should be DropChannels for channel 2.
        # (The order can be mixed up)
        # The remaining block should be FlowControl for 'Goodbye'.
        self.assertEqual(5, len(control_blocks))
        expected_opcode_and_flag = ((mux._MUX_OPCODE_DROP_CHANNEL << 5) |
                                    (1 << 4))
        self.assertTrue((expected_opcode_and_flag ==
                        (ord(control_blocks[3][0]) & 0xf0)) or
                        (expected_opcode_and_flag ==
                        (ord(control_blocks[4][0]) & 0xf0)))

    def test_fragmented_control_message(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(mux._INITIAL_NUMBER_OF_CHANNEL_SLOTS,
                                      mux._INITIAL_QUOTA_FOR_CLIENT)

        encoded_handshake = _create_request_header(path='/ping')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)

        # Replenish total 5 bytes in 3 FlowControls.
        flow_control = mux._create_flow_control(channel_id=2,
                                                replenished_quota=1,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        flow_control = mux._create_flow_control(channel_id=2,
                                                replenished_quota=2,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        flow_control = mux._create_flow_control(channel_id=2,
                                                replenished_quota=2,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))

        mux_handler.wait_until_done(timeout=2)

        messages = request.connection.get_written_control_messages(2)
        self.assertEqual(common.OPCODE_PING, messages[0]['opcode'])
        self.assertEqual('Ping!', messages[0]['message'])

    def test_channel_slot_violation_by_client(self):
        request = _create_mock_request()
        dispatcher = _MuxMockDispatcher()
        mux_handler = mux._MuxHandler(request, dispatcher)
        mux_handler.start()
        mux_handler.add_channel_slots(slots=1,
                                      send_quota=mux._INITIAL_QUOTA_FOR_CLIENT)

        encoded_handshake = _create_request_header(path='/echo')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=2, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)
        flow_control = mux._create_flow_control(channel_id=2,
                                                replenished_quota=10,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=2, message='Hello'))

        # This request should be rejected.
        encoded_handshake = _create_request_header(path='/echo')
        add_channel_request = _create_add_channel_request_frame(
            channel_id=3, encoding=0,
            encoded_handshake=encoded_handshake)
        request.connection.put_bytes(add_channel_request)
        flow_control = mux._create_flow_control(channel_id=3,
                                                replenished_quota=5,
                                                outer_frame_mask=True)
        request.connection.put_bytes(flow_control)

        request.connection.put_bytes(
            _create_logical_frame(channel_id=3, message='Hello'))

        request.connection.put_bytes(
            _create_logical_frame(channel_id=1, message='Goodbye'))
        request.connection.put_bytes(
            _create_logical_frame(channel_id=2, message='Goodbye'))

        mux_handler.wait_until_done(timeout=2)

        self.assertEqual(['Hello'], dispatcher.channel_events[2].messages)
        self.assertFalse(dispatcher.channel_events.has_key(3))

if __name__ == '__main__':
    unittest.main()

# vi:sts=4 sw=4 et