@@ -83,7 +83,7 @@ func TestPostToApmServerDataCompressed(t *testing.T) {
8383 apmproxy .WithLogger (zaptest .NewLogger (t ).Sugar ()),
8484 )
8585 require .NoError (t , err )
86- require .NoError (t , apmClient .PostToApmServer (context . Background (), agentData ))
86+ require .NoError (t , apmClient .PostToApmServer (t . Context (), agentData ))
8787}
8888
8989func TestPostToApmServerDataNotCompressed (t * testing.T ) {
@@ -129,7 +129,7 @@ func TestPostToApmServerDataNotCompressed(t *testing.T) {
129129 apmproxy .WithLogger (zaptest .NewLogger (t ).Sugar ()),
130130 )
131131 require .NoError (t , err )
132- require .NoError (t , apmClient .PostToApmServer (context . Background (), agentData ))
132+ require .NoError (t , apmClient .PostToApmServer (t . Context (), agentData ))
133133}
134134
135135func TestGracePeriod (t * testing.T ) {
@@ -178,7 +178,7 @@ func TestSetHealthyTransport(t *testing.T) {
178178 apmproxy .WithLogger (zaptest .NewLogger (t ).Sugar ()),
179179 )
180180 require .NoError (t , err )
181- apmClient .UpdateStatus (context . Background (), apmproxy .Healthy )
181+ apmClient .UpdateStatus (t . Context (), apmproxy .Healthy )
182182 assert .Equal (t , apmproxy .Healthy , apmClient .Status )
183183 assert .Equal (t , apmClient .ReconnectionCount , - 1 )
184184}
@@ -192,7 +192,7 @@ func TestSetFailingTransport(t *testing.T) {
192192 )
193193 require .NoError (t , err )
194194 apmClient .ReconnectionCount = 0
195- apmClient .UpdateStatus (context . Background (), apmproxy .Failing )
195+ apmClient .UpdateStatus (t . Context (), apmproxy .Failing )
196196 assert .Equal (t , apmproxy .Failing , apmClient .Status )
197197 assert .Equal (t , 1 , apmClient .ReconnectionCount )
198198}
@@ -203,8 +203,8 @@ func TestSetPendingTransport(t *testing.T) {
203203 apmproxy .WithLogger (zaptest .NewLogger (t ).Sugar ()),
204204 )
205205 require .NoError (t , err )
206- apmClient .UpdateStatus (context . Background (), apmproxy .Healthy )
207- apmClient .UpdateStatus (context . Background (), apmproxy .Failing )
206+ apmClient .UpdateStatus (t . Context (), apmproxy .Healthy )
207+ apmClient .UpdateStatus (t . Context (), apmproxy .Failing )
208208 require .Eventually (t , func () bool {
209209 return ! apmClient .IsUnhealthy ()
210210 }, 7 * time .Second , 50 * time .Millisecond )
@@ -218,8 +218,8 @@ func TestSetPendingTransportExplicitly(t *testing.T) {
218218 apmproxy .WithLogger (zaptest .NewLogger (t ).Sugar ()),
219219 )
220220 require .NoError (t , err )
221- apmClient .UpdateStatus (context . Background (), apmproxy .Healthy )
222- apmClient .UpdateStatus (context . Background (), apmproxy .Started )
221+ apmClient .UpdateStatus (t . Context (), apmproxy .Healthy )
222+ apmClient .UpdateStatus (t . Context (), apmproxy .Started )
223223 assert .Equal (t , apmproxy .Healthy , apmClient .Status )
224224 assert .Equal (t , apmClient .ReconnectionCount , - 1 )
225225}
@@ -230,8 +230,8 @@ func TestSetInvalidTransport(t *testing.T) {
230230 apmproxy .WithLogger (zaptest .NewLogger (t ).Sugar ()),
231231 )
232232 require .NoError (t , err )
233- apmClient .UpdateStatus (context . Background (), apmproxy .Healthy )
234- apmClient .UpdateStatus (context . Background (), "Invalid" )
233+ apmClient .UpdateStatus (t . Context (), apmproxy .Healthy )
234+ apmClient .UpdateStatus (t . Context (), "Invalid" )
235235 assert .Equal (t , apmproxy .Healthy , apmClient .Status )
236236 assert .Equal (t , apmClient .ReconnectionCount , - 1 )
237237}
@@ -274,12 +274,12 @@ func TestEnterBackoffFromHealthy(t *testing.T) {
274274 apmproxy .WithLogger (zaptest .NewLogger (t ).Sugar ()),
275275 )
276276 require .NoError (t , err )
277- apmClient .UpdateStatus (context . Background (), apmproxy .Healthy )
277+ apmClient .UpdateStatus (t . Context (), apmproxy .Healthy )
278278
279279 // Close the APM server early so that POST requests fail and that backoff is enabled
280280 apmServer .Close ()
281281
282- if err := apmClient .PostToApmServer (context . Background (), agentData ); err != nil {
282+ if err := apmClient .PostToApmServer (t . Context (), agentData ); err != nil {
283283 return
284284 }
285285 // No way to know for sure if failing or pending (0 sec grace period)
@@ -329,14 +329,14 @@ func TestEnterBackoffFromFailing(t *testing.T) {
329329 )
330330 require .NoError (t , err )
331331
332- apmClient .UpdateStatus (context . Background (), apmproxy .Healthy )
333- apmClient .UpdateStatus (context . Background (), apmproxy .Failing )
332+ apmClient .UpdateStatus (t . Context (), apmproxy .Healthy )
333+ apmClient .UpdateStatus (t . Context (), apmproxy .Failing )
334334 require .Eventually (t , func () bool {
335335 return ! apmClient .IsUnhealthy ()
336336 }, 7 * time .Second , 50 * time .Millisecond )
337337 assert .Equal (t , apmproxy .Started , apmClient .Status )
338338
339- require .Error (t , apmClient .PostToApmServer (context . Background (), agentData ))
339+ require .Error (t , apmClient .PostToApmServer (t . Context (), agentData ))
340340 assert .Equal (t , apmproxy .Failing , apmClient .Status )
341341 assert .Equal (t , 1 , apmClient .ReconnectionCount )
342342}
@@ -383,14 +383,14 @@ func TestAPMServerRecovery(t *testing.T) {
383383 )
384384 require .NoError (t , err )
385385
386- apmClient .UpdateStatus (context . Background (), apmproxy .Healthy )
387- apmClient .UpdateStatus (context . Background (), apmproxy .Failing )
386+ apmClient .UpdateStatus (t . Context (), apmproxy .Healthy )
387+ apmClient .UpdateStatus (t . Context (), apmproxy .Failing )
388388 require .Eventually (t , func () bool {
389389 return ! apmClient .IsUnhealthy ()
390390 }, 7 * time .Second , 50 * time .Millisecond )
391391 assert .Equal (t , apmproxy .Started , apmClient .Status )
392392
393- require .NoError (t , apmClient .PostToApmServer (context . Background (), agentData ))
393+ require .NoError (t , apmClient .PostToApmServer (t . Context (), agentData ))
394394 assert .Equal (t , apmproxy .Healthy , apmClient .Status )
395395}
396396
@@ -428,13 +428,13 @@ func TestAPMServerAuthFails(t *testing.T) {
428428 apmproxy .WithLogger (zaptest .NewLogger (t ).Sugar ()),
429429 )
430430 require .NoError (t , err )
431- apmClient .UpdateStatus (context . Background (), apmproxy .Healthy )
432- apmClient .UpdateStatus (context . Background (), apmproxy .Failing )
431+ apmClient .UpdateStatus (t . Context (), apmproxy .Healthy )
432+ apmClient .UpdateStatus (t . Context (), apmproxy .Failing )
433433 require .Eventually (t , func () bool {
434434 return ! apmClient .IsUnhealthy ()
435435 }, 7 * time .Second , 50 * time .Millisecond )
436436 assert .Equal (t , apmproxy .Started , apmClient .Status )
437- require .NoError (t , apmClient .PostToApmServer (context . Background (), agentData ))
437+ require .NoError (t , apmClient .PostToApmServer (t . Context (), agentData ))
438438 assert .NotEqual (t , apmproxy .Healthy , apmClient .Status )
439439}
440440
@@ -482,11 +482,11 @@ func TestAPMServerRatelimit(t *testing.T) {
482482 assert .Equal (t , apmproxy .Started , apmClient .Status )
483483
484484 // First request fails but does not trigger the backoff
485- require .NoError (t , apmClient .PostToApmServer (context . Background (), agentData ))
485+ require .NoError (t , apmClient .PostToApmServer (t . Context (), agentData ))
486486 assert .Equal (t , apmproxy .RateLimited , apmClient .Status )
487487
488488 // Followup request is successful
489- require .NoError (t , apmClient .PostToApmServer (context . Background (), agentData ))
489+ require .NoError (t , apmClient .PostToApmServer (t . Context (), agentData ))
490490 assert .Equal (t , apmproxy .Healthy , apmClient .Status )
491491}
492492
@@ -534,11 +534,11 @@ func TestAPMServerClientFail(t *testing.T) {
534534 assert .Equal (t , apmproxy .Started , apmClient .Status )
535535
536536 // First request fails but does not trigger the backoff
537- require .NoError (t , apmClient .PostToApmServer (context . Background (), agentData ))
537+ require .NoError (t , apmClient .PostToApmServer (t . Context (), agentData ))
538538 assert .Equal (t , apmproxy .ClientFailing , apmClient .Status )
539539
540540 // Followup request is successful
541- require .NoError (t , apmClient .PostToApmServer (context . Background (), agentData ))
541+ require .NoError (t , apmClient .PostToApmServer (t . Context (), agentData ))
542542 assert .Equal (t , apmproxy .Healthy , apmClient .Status )
543543}
544544
@@ -581,13 +581,13 @@ func TestContinuedAPMServerFailure(t *testing.T) {
581581 apmproxy .WithLogger (zaptest .NewLogger (t ).Sugar ()),
582582 )
583583 require .NoError (t , err )
584- apmClient .UpdateStatus (context . Background (), apmproxy .Healthy )
585- apmClient .UpdateStatus (context . Background (), apmproxy .Failing )
584+ apmClient .UpdateStatus (t . Context (), apmproxy .Healthy )
585+ apmClient .UpdateStatus (t . Context (), apmproxy .Failing )
586586 require .Eventually (t , func () bool {
587587 return ! apmClient .IsUnhealthy ()
588588 }, 7 * time .Second , 50 * time .Millisecond )
589589 assert .Equal (t , apmproxy .Started , apmClient .Status )
590- require .Error (t , apmClient .PostToApmServer (context . Background (), agentData ))
590+ require .Error (t , apmClient .PostToApmServer (t . Context (), agentData ))
591591 assert .Equal (t , apmproxy .Failing , apmClient .Status )
592592}
593593
@@ -628,7 +628,7 @@ func TestForwardApmData(t *testing.T) {
628628 require .NoError (t , err )
629629
630630 // Start forwarding APM data
631- ctx , cancel := context .WithCancel (context . Background ())
631+ ctx , cancel := context .WithCancel (t . Context ())
632632 var wg sync.WaitGroup
633633 wg .Add (1 )
634634 go func () {
@@ -703,7 +703,7 @@ func BenchmarkFlushAPMData(b *testing.B) {
703703 for j := 0 ; j < 99 ; j ++ {
704704 apmClient .LambdaDataChannel <- []byte (`{"log":{"message":this is test log"}}` )
705705 }
706- apmClient .FlushAPMData (context . Background ())
706+ apmClient .FlushAPMData (b . Context ())
707707 }
708708}
709709
@@ -742,7 +742,7 @@ func BenchmarkPostToAPM(b *testing.B) {
742742 b .ReportAllocs ()
743743 b .ResetTimer ()
744744 for i := 0 ; i < b .N ; i ++ {
745- if err := apmClient .PostToApmServer (context . Background (), agentData ); err != nil {
745+ if err := apmClient .PostToApmServer (b . Context (), agentData ); err != nil {
746746 b .Fatal (err )
747747 }
748748 }
0 commit comments