package io.digdag.core.agent;

import com.google.common.base.Optional;
import io.digdag.core.TempFileManager;
import io.digdag.spi.StorageObject;
import io.digdag.spi.TaskRequest;
import org.apache.commons.compress.archivers.tar.TarArchiveOutputStream;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.StandardOpenOption;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.zip.GZIPOutputStream;

import static org.hamcrest.core.Is.is;
import static org.junit.Assert.*;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

public class ExtractArchiveWorkspaceManagerTest
{
    @Rule
    public TemporaryFolder temporaryFolder = new TemporaryFolder();
    private ExtractArchiveWorkspaceManager extractArchiveWorkspaceManager;
    private byte[] archived;
    private TaskRequest taskRequest;
    private StorageObject storageObject;
    private StorageObject wrongStorageObject;

    @Before
    public void setUp()
            throws IOException
    {
        extractArchiveWorkspaceManager = new ExtractArchiveWorkspaceManager(new TempFileManager(temporaryFolder.getRoot().toPath()));

        File digFile = temporaryFolder.newFile();

        String wf = "+task:\n" +
                "  echo>: hello\n" +
                "";

        Files.write(digFile.toPath(), wf.getBytes(), StandardOpenOption.CREATE);

        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        GZIPOutputStream gzipOutputStream = new GZIPOutputStream(outputStream);
        TarArchiveOutputStream tarArchiveOutputStream = new TarArchiveOutputStream(gzipOutputStream);
        tarArchiveOutputStream.createArchiveEntry(digFile, "mydig.dig");
        tarArchiveOutputStream.close();

        archived = outputStream.toByteArray();

        taskRequest = mock(TaskRequest.class);
        when(taskRequest.getTaskName()).thenReturn("zzz");

        storageObject = mock(StorageObject.class);
        when(storageObject.getContentLength()).thenReturn((long) archived.length);
        when(storageObject.getContentInputStream()).thenReturn(new ByteArrayInputStream(archived));

        byte[] wrongArchive = "This isn't a proper archive content".getBytes();
        wrongStorageObject = mock(StorageObject.class);
        when(wrongStorageObject.getContentLength()).thenReturn((long) wrongArchive.length);
        when(wrongStorageObject.getContentInputStream()).thenReturn(new ByteArrayInputStream(wrongArchive));
    }

    @Test
    public void withExtractedArchive()
            throws Exception
    {
        WorkspaceManager.ArchiveProvider archiveProvider = mock(WorkspaceManager.ArchiveProvider.class);

        when(archiveProvider.open()).thenReturn(Optional.of(storageObject));

        AtomicInteger funcCounter = new AtomicInteger();
        extractArchiveWorkspaceManager.withExtractedArchive(taskRequest, archiveProvider, (path) -> funcCounter.incrementAndGet());
        assertThat(funcCounter.get(), is(1));

        assertEmptyWorkspace();
    }

    @Test
    public void withExtractedArchiveWithRetry()
            throws Exception
    {
        WorkspaceManager.ArchiveProvider archiveProvider = mock(WorkspaceManager.ArchiveProvider.class);

        when(archiveProvider.open()).thenReturn(Optional.of(wrongStorageObject),
                Optional.of(wrongStorageObject),
                Optional.of(storageObject));

        AtomicInteger funcCounter = new AtomicInteger();
        extractArchiveWorkspaceManager.withExtractedArchive(taskRequest, archiveProvider, (path) -> funcCounter.incrementAndGet());
        verify(archiveProvider, times(3)).open();
        assertThat(funcCounter.get(), is(1));

        assertEmptyWorkspace();
    }

    private void assertEmptyWorkspace()
    {
        assertThat(temporaryFolder.getRoot().toPath().resolve("workspace").toFile().listFiles().length, is(0));
    }
}
@@ -3220,32 +3220,82 @@
 ementAndGet());%0A
+        verify(archiveProvider, times(1)).open();%0A
         assertTh
@@ -4060,24 +4060,611 @@
 e();%0A    %7D%0A%0A
+    @Test%0A    public void withoutExtractedArchive()%0A            throws Exception%0A    %7B%0A        WorkspaceManager.ArchiveProvider archiveProvider = mock(WorkspaceManager.ArchiveProvider.class);%0A%0A        when(archiveProvider.open()).thenReturn(Optional.absent());%0A%0A        AtomicInteger funcCounter = new AtomicInteger();%0A        extractArchiveWorkspaceManager.withExtractedArchive(taskRequest, archiveProvider, (path) -%3E funcCounter.incrementAndGet());%0A        verify(archiveProvider, times(1)).open();%0A        assertThat(funcCounter.get(), is(1));%0A%0A        assertEmptyWorkspace();%0A    %7D%0A%0A
     private 

Add a few tests