|
4 | 4 | import static org.mockito.Mockito.*; |
5 | 5 |
|
6 | 6 | import com.databricks.jdbc.common.util.DatabricksThreadContextHolder; |
| 7 | +import com.databricks.jdbc.model.telemetry.StatementTelemetryDetails; |
7 | 8 | import com.databricks.jdbc.model.telemetry.latency.ChunkDetails; |
8 | 9 | import com.databricks.jdbc.model.telemetry.latency.OperationType; |
| 10 | +import com.databricks.jdbc.telemetry.TelemetryHelper; |
9 | 11 | import org.junit.jupiter.api.AfterEach; |
10 | 12 | import org.junit.jupiter.api.BeforeEach; |
11 | 13 | import org.junit.jupiter.api.Test; |
12 | 14 | import org.junit.jupiter.params.ParameterizedTest; |
13 | 15 | import org.junit.jupiter.params.provider.CsvSource; |
14 | 16 | import org.junit.jupiter.params.provider.EnumSource; |
| 17 | +import org.mockito.MockedStatic; |
15 | 18 |
|
16 | 19 | public class TelemetryCollectorTest { |
17 | 20 | private final TelemetryCollector handler = TelemetryCollector.getInstance(); |
@@ -71,16 +74,12 @@ void testRecordOperationLatency_WithCloseOperation() { |
71 | 74 | String methodName = "closeStatement"; |
72 | 75 | long latency = 100L; |
73 | 76 |
|
74 | | - // Ensure telemetry is collected for the statement id prior to close |
75 | | - handler.recordGetOperationStatus(TEST_STATEMENT_ID, 1L); |
76 | | - handler.recordOperationLatency(latency, methodName); |
77 | | - // Allow asynchronous export to complete and internal map to be cleared |
78 | | - try { |
79 | | - Thread.sleep(200); |
80 | | - } catch (InterruptedException e) { |
81 | | - Thread.currentThread().interrupt(); |
| 77 | + try (MockedStatic<TelemetryHelper> mockedStatic = mockStatic(TelemetryHelper.class)) { |
| 78 | + handler.recordOperationLatency(latency, methodName); |
| 79 | + |
| 80 | + mockedStatic.verify( |
| 81 | + () -> TelemetryHelper.exportTelemetryLog(any(StatementTelemetryDetails.class))); |
82 | 82 | } |
83 | | - assertNull(handler.getTelemetryDetails(TEST_STATEMENT_ID)); |
84 | 83 | } |
85 | 84 |
|
86 | 85 | @ParameterizedTest |
|
0 commit comments