Examples with TestParams org.opensearch.persistent.TestPersistentTasksPlugin.TestParams used on opensource projects

Example 21 with TestParams

use of org.opensearch.persistent.TestPersistentTasksPlugin.TestParams in project OpenSearch by opensearch-project.

the class PersistentTasksNodeServiceTests method testTaskCancellation.

public void testTaskCancellation() {
    AtomicLong capturedTaskId = new AtomicLong();
    AtomicReference<ActionListener<CancelTasksResponse>> capturedListener = new AtomicReference<>();
    Client client = mock(Client.class);
    when(client.settings()).thenReturn(Settings.EMPTY);
    PersistentTasksService persistentTasksService = new PersistentTasksService(null, null, client) {

        @Override
        void sendCancelRequest(final long taskId, final String reason, final ActionListener<CancelTasksResponse> listener) {
            capturedTaskId.set(taskId);
            capturedListener.set(listener);
        }

        @Override
        public void sendCompletionRequest(final String taskId, final long taskAllocationId, final Exception taskFailure, final ActionListener<PersistentTask<?>> listener) {
            fail("Shouldn't be called during Cluster State cancellation");
        }
    };
    @SuppressWarnings("unchecked") PersistentTasksExecutor<TestParams> action = mock(PersistentTasksExecutor.class);
    when(action.getExecutor()).thenReturn(ThreadPool.Names.SAME);
    when(action.getTaskName()).thenReturn("test");
    when(action.createTask(anyLong(), anyString(), anyString(), any(), any(), any())).thenReturn(new TestPersistentTasksPlugin.TestTask(1, "persistent", "test", "", new TaskId("cluster", 1), Collections.emptyMap()));
    PersistentTasksExecutorRegistry registry = new PersistentTasksExecutorRegistry(Collections.singletonList(action));
    int nonLocalNodesCount = randomInt(10);
    MockExecutor executor = new MockExecutor();
    TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet());
    PersistentTasksNodeService coordinator = new PersistentTasksNodeService(persistentTasksService, registry, taskManager, executor);
    ClusterState state = createInitialClusterState(nonLocalNodesCount, Settings.EMPTY);
    ClusterState newClusterState = state;
    // Allocate first task
    state = newClusterState;
    newClusterState = addTask(state, "test", null, "this_node");
    coordinator.clusterChanged(new ClusterChangedEvent("test", newClusterState, state));
    // Check the task is know to the task manager
    assertThat(taskManager.getTasks().size(), equalTo(1));
    AllocatedPersistentTask runningTask = (AllocatedPersistentTask) taskManager.getTasks().values().iterator().next();
    String persistentId = runningTask.getPersistentTaskId();
    long localId = runningTask.getId();
    // Make sure it returns correct status
    Task.Status status = runningTask.getStatus();
    assertThat(status.toString(), equalTo("{\"state\":\"STARTED\"}"));
    state = newClusterState;
    // Relocate the task to some other node or remove it completely
    if (randomBoolean()) {
        newClusterState = reallocateTask(state, persistentId, "some_other_node");
    } else {
        newClusterState = removeTask(state, persistentId);
    }
    coordinator.clusterChanged(new ClusterChangedEvent("test", newClusterState, state));
    // Make sure it returns correct status
    assertThat(taskManager.getTasks().size(), equalTo(1));
    assertThat(taskManager.getTasks().values().iterator().next().getStatus().toString(), equalTo("{\"state\":\"PENDING_CANCEL\"}"));
    // That should trigger cancellation request
    assertThat(capturedTaskId.get(), equalTo(localId));
    // Notify successful cancellation
    capturedListener.get().onResponse(new CancelTasksResponse(Collections.emptyList(), Collections.emptyList(), Collections.emptyList()));
    // finish or fail task
    if (randomBoolean()) {
        executor.get(0).task.markAsCompleted();
    } else {
        executor.get(0).task.markAsFailed(new IOException("test"));
    }
    // Check the task is now removed from task manager
    assertThat(taskManager.getTasks().values(), empty());
}