Skip to content

Commit

Permalink
Remove unneeded argument captors
Browse files Browse the repository at this point in the history 
Signed-off-by: Daniel Widdis <[email protected]>
  • Loading branch information
@dbwiddis
dbwiddis committed Dec 1, 2023
1 parent 8a81c60 commit ba6657b
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 23 deletions.
16 changes: 4 additions & 12 deletions src/test/java/org/opensearch/flowframework/workflow/CreateConnectorStepTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;

import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

Expand Down Expand Up @@ -81,15 +80,12 @@ public void testCreateConnector() throws IOException, ExecutionException, Interr
String connectorId = "connect";
CreateConnectorStep createConnectorStep = new CreateConnectorStep(machineLearningNodeClient, flowFrameworkIndicesHandler);

@SuppressWarnings("unchecked")
ArgumentCaptor<ActionListener<MLCreateConnectorResponse>> actionListenerCaptor = ArgumentCaptor.forClass(ActionListener.class);

doAnswer(invocation -> {
ActionListener<MLCreateConnectorResponse> actionListener = invocation.getArgument(1);
MLCreateConnectorResponse output = new MLCreateConnectorResponse(connectorId);
actionListener.onResponse(output);
return null;
}).when(machineLearningNodeClient).createConnector(any(MLCreateConnectorInput.class), actionListenerCaptor.capture());
}).when(machineLearningNodeClient).createConnector(any(MLCreateConnectorInput.class), any());

CompletableFuture<WorkflowData> future = createConnectorStep.execute(
inputData.getNodeId(),
Expand All @@ -98,8 +94,7 @@ public void testCreateConnector() throws IOException, ExecutionException, Interr
Collections.emptyMap()
);

verify(machineLearningNodeClient).createConnector(any(MLCreateConnectorInput.class), actionListenerCaptor.capture());

verify(machineLearningNodeClient).createConnector(any(MLCreateConnectorInput.class), any());
assertTrue(future.isDone());
assertEquals(connectorId, future.get().getContent().get("connector_id"));

Expand All @@ -108,14 +103,11 @@ public void testCreateConnector() throws IOException, ExecutionException, Interr
public void testCreateConnectorFailure() throws IOException {
CreateConnectorStep createConnectorStep = new CreateConnectorStep(machineLearningNodeClient, flowFrameworkIndicesHandler);

@SuppressWarnings("unchecked")
ArgumentCaptor<ActionListener<MLCreateConnectorResponse>> actionListenerCaptor = ArgumentCaptor.forClass(ActionListener.class);

doAnswer(invocation -> {
ActionListener<MLCreateConnectorResponse> actionListener = invocation.getArgument(1);
actionListener.onFailure(new FlowFrameworkException("Failed to create connector", RestStatus.INTERNAL_SERVER_ERROR));
return null;
}).when(machineLearningNodeClient).createConnector(any(MLCreateConnectorInput.class), actionListenerCaptor.capture());
}).when(machineLearningNodeClient).createConnector(any(MLCreateConnectorInput.class), any());

CompletableFuture<WorkflowData> future = createConnectorStep.execute(
inputData.getNodeId(),
Expand All @@ -124,7 +116,7 @@ public void testCreateConnectorFailure() throws IOException {
Collections.emptyMap()
);

verify(machineLearningNodeClient).createConnector(any(MLCreateConnectorInput.class), actionListenerCaptor.capture());
verify(machineLearningNodeClient).createConnector(any(MLCreateConnectorInput.class), any());

assertTrue(future.isCompletedExceptionally());
ExecutionException ex = assertThrows(ExecutionException.class, () -> future.get().getContent());
Expand Down
15 changes: 4 additions & 11 deletions src/test/java/org/opensearch/flowframework/workflow/DeleteConnectorStepTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;

import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

Expand Down Expand Up @@ -53,25 +52,22 @@ public void testDeleteConnector() throws IOException, ExecutionException, Interr
String connectorId = randomAlphaOfLength(5);
DeleteConnectorStep deleteConnectorStep = new DeleteConnectorStep(machineLearningNodeClient);

@SuppressWarnings("unchecked")
ArgumentCaptor<ActionListener<DeleteResponse>> actionListenerCaptor = ArgumentCaptor.forClass(ActionListener.class);

doAnswer(invocation -> {
String connectorIdArg = invocation.getArgument(0);
ActionListener<DeleteResponse> actionListener = invocation.getArgument(1);
ShardId shardId = new ShardId(new Index("indexName", "uuid"), 1);
DeleteResponse output = new DeleteResponse(shardId, connectorIdArg, 1, 1, 1, true);
actionListener.onResponse(output);
return null;
}).when(machineLearningNodeClient).deleteConnector(any(String.class), actionListenerCaptor.capture());
}).when(machineLearningNodeClient).deleteConnector(any(String.class), any());

CompletableFuture<WorkflowData> future = deleteConnectorStep.execute(
inputData.getNodeId(),
inputData,
Map.of("step_1", new WorkflowData(Map.of("connector_id", connectorId), "workflowId", "nodeId")),
Map.of("step_1", "connector_id")
);
verify(machineLearningNodeClient).deleteConnector(any(String.class), actionListenerCaptor.capture());
verify(machineLearningNodeClient).deleteConnector(any(String.class), any());

assertTrue(future.isDone());
assertEquals(connectorId, future.get().getContent().get("connector_id"));
Expand All @@ -96,14 +92,11 @@ public void testNoConnectorIdInOutput() throws IOException {
public void testDeleteConnectorFailure() throws IOException {
DeleteConnectorStep deleteConnectorStep = new DeleteConnectorStep(machineLearningNodeClient);

@SuppressWarnings("unchecked")
ArgumentCaptor<ActionListener<DeleteResponse>> actionListenerCaptor = ArgumentCaptor.forClass(ActionListener.class);

doAnswer(invocation -> {
ActionListener<MLCreateConnectorResponse> actionListener = invocation.getArgument(1);
actionListener.onFailure(new FlowFrameworkException("Failed to delete connector", RestStatus.INTERNAL_SERVER_ERROR));
return null;
}).when(machineLearningNodeClient).deleteConnector(any(String.class), actionListenerCaptor.capture());
}).when(machineLearningNodeClient).deleteConnector(any(String.class), any());

CompletableFuture<WorkflowData> future = deleteConnectorStep.execute(
inputData.getNodeId(),
Expand All @@ -112,7 +105,7 @@ public void testDeleteConnectorFailure() throws IOException {
Map.of("step_1", "connector_id")
);

verify(machineLearningNodeClient).deleteConnector(any(String.class), actionListenerCaptor.capture());
verify(machineLearningNodeClient).deleteConnector(any(String.class), any());

assertTrue(future.isCompletedExceptionally());
ExecutionException ex = assertThrows(ExecutionException.class, () -> future.get().getContent());
Expand Down

0 comments on commit ba6657b

Please sign in to comment.