|
30 | 30 | import com.datastax.oss.driver.api.core.config.DriverExecutionProfile; |
31 | 31 | import com.datastax.oss.driver.api.core.cql.AsyncResultSet; |
32 | 32 | import com.datastax.oss.driver.api.core.servererrors.BootstrappingException; |
| 33 | +import com.datastax.oss.driver.api.core.session.Request; |
| 34 | +import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; |
33 | 35 | import com.datastax.oss.driver.api.core.tracker.RequestTracker; |
34 | 36 | import com.datastax.oss.driver.internal.core.tracker.NoopRequestTracker; |
35 | 37 | import com.datastax.oss.protocol.internal.ProtocolConstants; |
@@ -92,6 +94,61 @@ public void should_invoke_request_tracker() { |
92 | 94 | } |
93 | 95 | } |
94 | 96 |
|
| 97 | + @Test |
| 98 | + public void should_pass_node_specific_log_prefix_to_node_success_tracker() { |
| 99 | + RequestIdGenerator requestIdGenerator = |
| 100 | + new RequestIdGenerator() { |
| 101 | + @Override |
| 102 | + public String getSessionRequestId() { |
| 103 | + return "session-id"; |
| 104 | + } |
| 105 | + |
| 106 | + @Override |
| 107 | + public String getNodeRequestId(Request request, String parentId) { |
| 108 | + return "node-id"; |
| 109 | + } |
| 110 | + }; |
| 111 | + |
| 112 | + try (RequestHandlerTestHarness harness = |
| 113 | + RequestHandlerTestHarness.builder() |
| 114 | + .withDefaultIdempotence(true) |
| 115 | + .withRequestIdGenerator(requestIdGenerator) |
| 116 | + .withResponse(node1, defaultFrameOf(singleRow())) |
| 117 | + .build()) { |
| 118 | + |
| 119 | + RequestTracker requestTracker = mock(RequestTracker.class); |
| 120 | + when(harness.getContext().getRequestTracker()).thenReturn(requestTracker); |
| 121 | + |
| 122 | + CompletionStage<AsyncResultSet> resultSetFuture = |
| 123 | + new CqlRequestHandler( |
| 124 | + UNDEFINED_IDEMPOTENCE_STATEMENT, |
| 125 | + harness.getSession(), |
| 126 | + harness.getContext(), |
| 127 | + "test") |
| 128 | + .handle(); |
| 129 | + |
| 130 | + assertThatStage(resultSetFuture) |
| 131 | + .isSuccess( |
| 132 | + resultSet -> { |
| 133 | + verify(requestTracker) |
| 134 | + .onNodeSuccess( |
| 135 | + any(), |
| 136 | + anyLong(), |
| 137 | + any(DriverExecutionProfile.class), |
| 138 | + eq(node1), |
| 139 | + eq("test|node-id|0")); |
| 140 | + verify(requestTracker) |
| 141 | + .onSuccess( |
| 142 | + any(), |
| 143 | + anyLong(), |
| 144 | + any(DriverExecutionProfile.class), |
| 145 | + eq(node1), |
| 146 | + eq("test|session-id")); |
| 147 | + verifyNoMoreInteractions(requestTracker); |
| 148 | + }); |
| 149 | + } |
| 150 | + } |
| 151 | + |
95 | 152 | @Test |
96 | 153 | public void should_not_invoke_noop_request_tracker() { |
97 | 154 | try (RequestHandlerTestHarness harness = |
|
0 commit comments