daemon: fix flaky test

Signed-off-by: Denys Smirnov <[email protected]>

daemon/pool.go

@@ -614,6 +614,13 @@ func (dp *DriverPool) manageDrivers() {
 // pool scaling limit. The caller should put the instance back to the pool even if the
 // driver fails.
 func (dp *DriverPool) getIdle(rctx context.Context) (Driver, error) {
+	// don't do anything if the request is already cancelled,
+	// or we will have to "rollback" it later
+	select {
+	case <-rctx.Done():
+		return nil, rctx.Err()
+	default:
+	}
 	// fast path - get an idle driver directly from the pool
 	// this function executes on the current goroutine
 	if d, ok := dp.peekIdle(); ok {

daemon/pool_test.go

@@ -2,6 +2,7 @@ package daemon
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"runtime"
 	"sync"
@@ -28,7 +29,9 @@ func TestDriverPoolClose_StartNoopClose(t *testing.T) {
 	err = dp.Stop()
 	require.True(ErrPoolClosed.Is(err), "%v", err)
 
-	err = dp.ExecuteCtx(ctx, nil)
+	err = dp.ExecuteCtx(ctx, func(ctx context.Context, d Driver) error {
+		return errors.New("should not happen")
+	})
 	require.True(ErrPoolClosed.Is(err), "%v", err)
 }
 
@@ -65,7 +68,9 @@ func TestDriverPoolExecute_Timeout(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond)
 	defer cancel()
 
-	err = dp.ExecuteCtx(ctx, nil)
+	err = dp.ExecuteCtx(ctx, func(ctx context.Context, d Driver) error {
+		return errors.New("should not happen")
+	})
 	require.True(err == context.DeadlineExceeded)
 }