@@ -172,7 +172,7 @@ func TestPartialFlush(t *testing.T) {
172
172
t .Setenv ("DD_TRACE_PARTIAL_FLUSH_ENABLED" , "true" )
173
173
t .Setenv ("DD_TRACE_PARTIAL_FLUSH_MIN_SPANS" , "2" )
174
174
t .Run ("WithFlush" , func (t * testing.T ) {
175
- telemetryClient := new (telemetrytest.MockClient )
175
+ telemetryClient := new (telemetrytest.RecordClient )
176
176
telemetryClient .ProductStarted (telemetry .NamespaceTracers )
177
177
defer testTelemetryClient (telemetryClient )()
178
178
tracer , transport , flush , stop := startTestTracer (t )
@@ -198,10 +198,9 @@ func TestPartialFlush(t *testing.T) {
198
198
comparePayloadSpans (t , children [0 ], ts [0 ][0 ])
199
199
comparePayloadSpans (t , children [1 ], ts [0 ][1 ])
200
200
201
- telemetryClient .AssertCalled (t , "Count" , telemetry .NamespaceTracers , "trace_partial_flush.count" , 1.0 , []string {"reason:large_trace" }, true )
202
- // TODO: (Support MetricKindDist) Re-enable these when we actually support `MetricKindDist`
203
- //telemetryClient.AssertCalled(t, "Record", telemetry.NamespaceTracers, "trace_partial_flush.spans_closed", 2.0, []string(nil), true) // Typed-nil here to not break usage of reflection in `mock` library.
204
- //telemetryClient.AssertCalled(t, "Record", telemetry.NamespaceTracers, "trace_partial_flush.spans_remaining", 1.0, []string(nil), true)
201
+ assert .Equal (t , 1.0 , telemetryClient .Count (telemetry .NamespaceTracers , "trace_partial_flush.count" , []string {"reason:large_trace" }).Get ())
202
+ assert .Equal (t , 2.0 , telemetryClient .Distribution (telemetry .NamespaceTracers , "trace_partial_flush.spans_closed" , nil ).Get ())
203
+ assert .Equal (t , 1.0 , telemetryClient .Distribution (telemetry .NamespaceTracers , "trace_partial_flush.spans_remaining" , nil ).Get ())
205
204
206
205
root .Finish ()
207
206
flush (1 )
@@ -214,9 +213,6 @@ func TestPartialFlush(t *testing.T) {
214
213
assert .Equal (t , 1.0 , ts [0 ][1 ].Metrics [keySamplingPriority ]) // the tag should only be on the first span in the chunk
215
214
comparePayloadSpans (t , root .(* span ), tsRoot [0 ][0 ])
216
215
comparePayloadSpans (t , children [2 ], tsRoot [0 ][1 ])
217
- telemetryClient .AssertNumberOfCalls (t , "Count" , 1 )
218
- // TODO: (Support MetricKindDist) Re-enable this when we actually support `MetricKindDist`
219
- // telemetryClient.AssertNumberOfCalls(t, "Record", 2)
220
216
})
221
217
222
218
// This test covers an issue where partial flushing + a rate sampler would panic
0 commit comments