From b0eb1b8c92644844952bb24d4537a72067a27572 Mon Sep 17 00:00:00 2001 From: cchenggit Date: Mon, 2 Dec 2024 21:45:41 +0800 Subject: [PATCH] fix(tests): fix flaky test (#67) --- pkg/pool/pool_test.go | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/pkg/pool/pool_test.go b/pkg/pool/pool_test.go index 7d1648e..a899dc6 100644 --- a/pkg/pool/pool_test.go +++ b/pkg/pool/pool_test.go @@ -62,19 +62,22 @@ func TestShutdown(t *testing.T) { } func TestGracefulShutdown(t *testing.T) { - var counter int64 - atomic.StoreInt64(&counter, 0) + var counter atomic.Int64 pool := NewPool(100, 100) + wait := sync.WaitGroup{} + wait.Add(100) for i := 0; i < 100; i++ { err := pool.SubmitFn(time.Second, func() { + wait.Done() time.Sleep(time.Second) - atomic.AddInt64(&counter, 1) + counter.Add(1) }) assert.NoError(t, err) } + wait.Wait() // wait for all tasks to be scheduled pool.Shutdown() - assert.EqualValues(t, 100, counter) // all submitted and scheduled tasks should be executed successfully + assert.EqualValues(t, 100, counter.Load()) // all submitted and scheduled tasks should be executed successfully }