package io.digdag.core.agent;

import com.google.common.io.Resources;
import io.digdag.client.DigdagClient;
import io.digdag.client.config.Config;
import io.digdag.client.config.ConfigFactory;
import io.digdag.client.config.ConfigUtils;
import io.digdag.core.Limits;
import io.digdag.core.workflow.OperatorTestingUtils;
import io.digdag.spi.SecretStoreManager;
import io.digdag.spi.TaskExecutionException;
import io.digdag.spi.TaskRequest;
import io.digdag.spi.TaskResult;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;

import java.io.IOException;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.junit.Assert.assertEquals;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class OperatorManagerTest
{
    private AgentConfig agentConfig = AgentConfig.defaultBuilder().build();
    private AgentId agentId = AgentId.of("dummy");
    @Mock TaskCallbackApi callback;
    private ConfigFactory cf = new ConfigFactory(DigdagClient.objectMapper());
    @Mock OperatorRegistry registry;
    @Mock SecretStoreManager secretStoreManager;
    private Limits limits = new Limits(cf.create());
    private Config simpleConfig = cf.fromJsonString("{\"echo>\":\"hello\"}");

    private OperatorManager operatorManager;

    @Before
    public void setUp()
    {
        ConfigEvalEngine evalEngine = new ConfigEvalEngine(ConfigUtils.newConfig());
        WorkspaceManager workspaceManager = new LocalWorkspaceManager();

        operatorManager = new OperatorManager(
                agentConfig, agentId, callback, workspaceManager, cf,
                evalEngine, registry, secretStoreManager, limits);
    }

    @Test
    public void testFilterConfigForLogging()
            throws IOException
    {
        Config src = cf.fromJsonString(Resources.toString(OperatorManagerTest.class.getResource("/io/digdag/core/agent/operator_manager/filter_config_src.json"), UTF_8));
        Config srcBackup = src.deepCopy();
        Config expectedConfig = cf.fromJsonString(Resources.toString(OperatorManagerTest.class.getResource("/io/digdag/core/agent/operator_manager/filter_config_expected.json"), UTF_8));

        Config filteredConfig = operatorManager.filterConfigForLogging(src);

        assertEquals(src, srcBackup); // src must not be modified
        assertEquals(expectedConfig.getKeys().size(), filteredConfig.getKeys().size());
        for (String k : expectedConfig.getKeys()) {
            assertEquals(expectedConfig.get(k, Object.class), filteredConfig.get(k, Object.class));
        }
    }

    @Test
    public void testRunWithHeartbeatWithSuccessTask()
    {
        TaskRequest taskRequest = OperatorTestingUtils.newTaskRequest(simpleConfig);

        TaskResult result = mock(TaskResult.class);
        OperatorManager om = spy(operatorManager);
        doReturn(result).when(om).callExecutor(any(), any(), any());
        om.runWithHeartbeat(taskRequest);
        verify(callback, times(1)).taskSucceeded(eq(taskRequest), any(), eq(result));
        verify(callback, times(0)).taskFailed(any(), any(), any());
    }

    @Test
    public void testRunWithHeartbeatWithFailedTask()
    {
        TaskRequest taskRequest = OperatorTestingUtils.newTaskRequest(simpleConfig);

        OperatorManager om = spy(operatorManager);
        doThrow(new TaskExecutionException("Zzz")).when(om).callExecutor(any(), any(), any());
        om.runWithHeartbeat(taskRequest);
        verify(callback, times(0)).taskSucceeded(any(), any(), any());
        verify(callback, times(1)).taskFailed(eq(taskRequest), any(), any());
    }
}

@@ -140,16 +140,62 @@
 Config;%0A
+import io.digdag.client.config.ConfigElement;%0A
 import i
@@ -384,138 +384,9 @@
 spi.
-SecretStoreManager;%0Aimport io.digdag.spi.TaskExecutionException;%0Aimport io.digdag.spi.TaskRequest;%0Aimport io.digdag.spi.TaskResult
+*
 ;%0Aim
@@ -3166,24 +3166,108 @@
 (), any());%0A
+        verify(callback, times(0)).retryTask(any(), any(), anyInt(), any(), any());%0A
     %7D%0A%0A    @
@@ -3749,20 +3749,769 @@
 (), any());%0A
+        verify(callback, times(0)).retryTask(any(), any(), anyInt(), any(), any());%0A    %7D%0A%0A    @Test%0A    public void testRunWithHeartbeatWithFailedTaskWithRetryableFailure()%0A    %7B%0A        TaskRequest taskRequest = OperatorTestingUtils.newTaskRequest(simpleConfig);%0A%0A        OperatorManager om = spy(operatorManager);%0A        doThrow(TaskExecutionException.ofNextPolling(42, ConfigElement.empty()))%0A                .when(om)%0A                .callExecutor(any(), any(), any());%0A        om.runWithHeartbeat(taskRequest);%0A        verify(callback, times(0)).taskSucceeded(any(), any(), any());%0A        verify(callback, times(0)).taskFailed(any(), any(), any());%0A        verify(callback, times(1)).retryTask(eq(taskRequest), any(), eq(42), any(), any());%0A
     %7D%0A%7D%0A

Add a test