package acceptance;

import io.digdag.client.DigdagClient;
import io.digdag.client.api.RestSessionAttempt;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import utils.TemporaryDigdagServer;
import utils.TestUtils;

import java.nio.file.Path;
import java.time.Duration;

import static org.hamcrest.Matchers.is;
import static org.junit.Assert.assertThat;
import static utils.TestUtils.addWorkflow;
import static utils.TestUtils.attemptFailure;
import static utils.TestUtils.expect;
import static utils.TestUtils.pushAndStart;

public class AttemptTimeoutIT
{
    @Rule
    public TemporaryFolder folder = new TemporaryFolder();

    @Rule
    public TemporaryDigdagServer server = TemporaryDigdagServer.builder()
            .configuration(
                    "executor.attempt_ttl = 5s",
                    "executor.attempt_reaping_interval = 1s"
            )
            .build();

    private Path projectDir;
    private DigdagClient client;

    @Before
    public void setUp()
            throws Exception
    {
        projectDir = folder.newFolder().toPath();

        client = DigdagClient.builder()
                .host(server.host())
                .port(server.port())
                .build();
    }

    @Test
    public void testAttemptTimeout()
            throws Exception
    {
        addWorkflow(projectDir, "acceptance/attempt_timeout/attempt_timeout.dig");
        long attemptId = pushAndStart(server.endpoint(), projectDir, "attempt_timeout");
        expect(Duration.ofSeconds(30), attemptFailure(server.endpoint(), attemptId));

        RestSessionAttempt attempt = client.getSessionAttempt(attemptId);
        assertThat(attempt.getCancelRequested(), is(true));
    }
}

@@ -250,32 +250,8 @@
 ver;
-%0Aimport utils.TestUtils;
 %0A%0Aim
@@ -433,54 +433,8 @@
 ow;%0A
-import static utils.TestUtils.attemptFailure;%0A
 impo
@@ -774,9 +774,10 @@
 l = 
-5
+10
 s%22,%0A
@@ -1464,224 +1464,542 @@
 %22);%0A
-        expect(Duration.ofSeconds(30), attemptFailure(server.endpoint(), attemptId));%0A%0A        RestSessionAttempt attempt = client.getSessionAttempt(attemptId);%0A        assertThat(attempt.getCancelRequested(), is(tru
+%0A        // Expect the attempt to get canceled%0A        expect(Duration.ofMinutes(1), () -%3E client.getSessionAttempt(attemptId).getCancelRequested());%0A%0A        // And then the attempt should be done pretty soon%0A        expect(Duration.ofMinutes(1), () -%3E client.getSessionAttempt(attemptId).getDone());%0A%0A        RestSessionAttempt attempt = client.getSessionAttempt(attemptId);%0A        assertThat(attempt.getDone(), is(true));%0A        assertThat(attempt.getCancelRequested(), is(true));%0A        assertThat(attempt.getSuccess(), is(fals
 e));

improve test