feat: [DGP-730] jitter TestAPI polling from 0.5 to 1.5 of the pollInterval (#401)

Author: bsalomon-snyk

pkg/apiclients/testapi/testapi.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math/rand"
 	"net/http"
 	"net/url"
 	"sync"
@@ -24,6 +25,7 @@ type config struct {
 	APIVersion            string
 	Logger                *zerolog.Logger
 	lowLevelClientOptions []ClientOption // Options for the oapi-codegen client
+	jitterFunc            func(time.Duration) time.Duration
 }
 
 // ConfigOption allows setting custom parameters during construction
@@ -85,6 +87,13 @@ func WithCustomRequestEditorFn(fn RequestEditorFn) ConfigOption {
 	}
 }
 
+// WithJitterFunc allows setting a custom jitter function for polling.
+func WithJitterFunc(fn func(time.Duration) time.Duration) ConfigOption {
+	return func(c *config) {
+		c.jitterFunc = fn
+	}
+}
+
 type client struct {
 	lowLevelClient ClientWithResponsesInterface
 	config         config
@@ -235,6 +244,7 @@ func NewTestClient(serverBaseUrl string, options ...ConfigOption) (TestClient, e
 	cfg := config{
 		PollInterval: DefaultPollInterval,
 		APIVersion:   DefaultAPIVersion,
+		jitterFunc:   Jitter,
 	}
 
 	for _, opt := range options {
@@ -441,7 +451,9 @@ func contextCanceledError(operationDescription string, contextError error) error
 
 // Query the job endpoint until we're redirected to its 'related' link containing results
 func (h *testHandle) pollJobToCompletion(ctx context.Context) (*uuid.UUID, error) {
-	ticker := time.NewTicker(h.client.config.PollInterval)
+	cfg := h.client.config
+
+	ticker := time.NewTicker(cfg.PollInterval)
 	defer ticker.Stop()
 
 	getJobParams := &GetJobParams{Version: h.client.config.APIVersion}
@@ -463,6 +475,7 @@ func (h *testHandle) pollJobToCompletion(ctx context.Context) (*uuid.UUID, error
 				if stopPolling {
 					return nil, jobErr
 				}
+				ticker.Reset(cfg.jitterFunc(cfg.PollInterval))
 				continue
 
 			case http.StatusSeeOther:
@@ -478,6 +491,7 @@ func (h *testHandle) pollJobToCompletion(ctx context.Context) (*uuid.UUID, error
 					Str("orgID", h.orgID.String()).
 					Str("jobID", h.jobID.String()).
 					Msg("Job polling returned 404 Not Found, continuing polling in case of delayed job creation")
+				ticker.Reset(cfg.jitterFunc(cfg.PollInterval))
 				continue
 
 			default:
@@ -691,3 +705,13 @@ func (r *testResult) handleFindingsError(err error, partialFindings []FindingDat
 func ptr[T any](v T) *T {
 	return &v
 }
+
+// Jitter returns a random duration between 0.5 and 1.5 of the given duration.
+func Jitter(d time.Duration) time.Duration {
+	if d <= 0 {
+		return d
+	}
+	minDur := int64(float64(d) * 0.5)
+	maxDur := int64(float64(d) * 1.5)
+	return time.Duration(rand.Int63n(maxDur-minDur) + minDur)
+}

pkg/apiclients/testapi/testapi_test.go

@@ -180,7 +180,7 @@ func Test_StartTest_Success(t *testing.T) {
 	// Create our test client
 	testHTTPClient := newTestHTTPClient(t, server)
 	testClient, err := testapi.NewTestClient(server.URL,
-		testapi.WithPollInterval(1*time.Second), // short poll interval for faster tests
+		testapi.WithPollInterval(1*time.Second),
 		testapi.WithCustomHTTPClient(testHTTPClient),
 	)
 	require.NoError(t, err)
@@ -1047,3 +1047,85 @@ func assertTestFinishedWithOutcomeErrorsAndWarnings(t *testing.T, result testapi
 		assert.Nil(t, result.GetWarnings())
 	}
 }
+
+// TestJitter verifies times returned are 0.5-1.5 times the original value,
+// and invalid times are return unmodified.
+func TestJitter(t *testing.T) {
+	t.Parallel()
+	t.Run("returns original duration for zero or negative input", func(t *testing.T) {
+		t.Parallel()
+		assert.Equal(t, time.Duration(0), testapi.Jitter(0))
+		assert.Equal(t, time.Duration(-1), testapi.Jitter(-1))
+	})
+
+	t.Run("returns duration within 0.5x to 1.5x of input", func(t *testing.T) {
+		t.Parallel()
+		duration := 100 * time.Millisecond
+		minDur := time.Duration(float64(duration) * 0.5)
+		maxDur := time.Duration(float64(duration) * 1.5)
+
+		for range 100 {
+			jittered := testapi.Jitter(duration)
+			assert.GreaterOrEqual(t, jittered, minDur)
+			assert.LessOrEqual(t, jittered, maxDur)
+		}
+	})
+}
+
+// Test_Wait_CallsJitter ensures Jitter is called while polling for job completion.
+func Test_Wait_CallsJitter(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
+	testData := setupTestScenario(t)
+
+	params := testapi.StartTestParams{
+		OrgID:   testData.OrgID.String(),
+		Subject: testData.TestSubjectCreate,
+	}
+
+	// Mock Jitter
+	var jitterCalled bool
+	jitterFunc := func(d time.Duration) time.Duration {
+		jitterCalled = true
+		return d
+	}
+
+	// Mock server handler
+	handlerConfig := TestAPIHandlerConfig{
+		OrgID:       testData.OrgID,
+		JobID:       testData.JobID,
+		TestID:      testData.TestID,
+		APIVersion:  testapi.DefaultAPIVersion,
+		PollCounter: testData.PollCounter,
+		JobPollResponses: []JobPollResponseConfig{
+			{Status: testapi.Pending}, // First poll
+			{ShouldRedirect: true},    // Second poll, redirects
+		},
+		FinalTestResult: FinalTestResultConfig{
+			Outcome: testapi.Pass,
+		},
+	}
+	handler := newTestAPIMockHandler(t, handlerConfig)
+	server, cleanup := startMockServer(t, handler)
+	defer cleanup()
+
+	// Act
+	testHTTPClient := newTestHTTPClient(t, server)
+	testClient, err := testapi.NewTestClient(server.URL,
+		testapi.WithPollInterval(1*time.Second),
+		testapi.WithCustomHTTPClient(testHTTPClient),
+		testapi.WithJitterFunc(jitterFunc),
+	)
+	require.NoError(t, err)
+
+	handle, err := testClient.StartTest(ctx, params)
+	require.NoError(t, err)
+	require.NotNil(t, handle)
+
+	err = handle.Wait(ctx)
+	require.NoError(t, err)
+
+	// Assert
+	assert.True(t, jitterCalled)
+}