Add unit test.

go.mod

@@ -7,8 +7,10 @@ toolchain go1.22.1
 require (
 	github.com/cenkalti/backoff/v4 v4.2.1
 	github.com/google/uuid v1.6.0
+	github.com/neilotoole/slogt v1.1.0
 	github.com/spf13/cobra v1.8.0
 	github.com/spf13/pflag v1.0.5
+	github.com/stretchr/testify v1.8.4
 	google.golang.org/grpc v1.63.2
 	google.golang.org/protobuf v1.33.0
 	k8s.io/apimachinery v0.30.0
@@ -17,11 +19,15 @@ require (
 )
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/go-logr/logr v1.4.1 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/kr/text v0.2.0 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	golang.org/x/net v0.23.0 // indirect
 	golang.org/x/sys v0.18.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240227224415-6ceb2ff114de // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

internal/metrics/submitter/submitter.go

@@ -17,6 +17,7 @@ type Submitter struct {
 	done    chan struct{}
 	client  gen.MetricsClient
 	logger  *slog.Logger
+	timer   backoff.Timer // for testing
 }
 
 // NewSubmitter creates a new submitter object.
@@ -38,11 +39,11 @@ func (s *Submitter) Chan() chan<- *gen.Result {
 }
 
 // Run accepts metrics on the channel and submits them to the client passed to constructor until Await is called.
-func (s *Submitter) Run(ctx context.Context) {
+func (s *Submitter) Run(ctx context.Context) (err error) {
 	defer func() { s.done <- struct{}{} }()
-	hostName, err := os.Hostname()
-	if err != nil {
-		s.logger.WarnContext(ctx, "could not obtain hostname", "error", err)
+	hostName, hostErr := os.Hostname()
+	if hostErr != nil {
+		s.logger.WarnContext(ctx, "could not obtain hostname", "error", hostErr)
 		hostName = "unknown"
 	}
 
@@ -53,33 +54,36 @@ func (s *Submitter) Run(ctx context.Context) {
 		metrics = append(metrics, metric)
 	}
 
-	if err = s.submit(ctx, metrics); err == nil {
-		s.logger.InfoContext(ctx, "metrics submitted")
-		return
-	}
-	s.logger.ErrorContext(ctx, "metric Submit RPC failed, retrying", "error", err)
-	b := backoff.NewExponentialBackOff()
-	b.InitialInterval = 10 * time.Second
-	b.MaxElapsedTime = 0
-	ticker := backoff.NewTicker(backoff.WithContext(b, ctx))
+	ticker := newTicker(ctx, s.timer)
 	defer ticker.Stop()
+
 	for {
 		select {
-		case <-ctx.Done():
-			if ctx.Err() != nil {
-				s.logger.ErrorContext(ctx, "giving up retrying metrics submission", "error", ctx.Err())
-			}
-			return
 		case <-ticker.C:
 			if err = s.submit(ctx, metrics); err == nil {
 				s.logger.InfoContext(ctx, "metrics submitted")
 				return
 			}
 			s.logger.ErrorContext(ctx, "metric Submit RPC failed, retrying", "error", err)
+		case <-ctx.Done():
+			if ctx.Err() != nil {
+				s.logger.ErrorContext(ctx, "giving up retrying metrics submission", "error", ctx.Err())
+				err = ctx.Err()
+			}
+			return
 		}
 	}
 }
 
+// newTicker returns a ticker that ticks once immediately, and then backs off exponentially forever.
+// Caller is responsible for calling Stop() on it eventually.
+func newTicker(ctx context.Context, timer backoff.Timer) *backoff.Ticker {
+	b := backoff.NewExponentialBackOff()
+	b.InitialInterval = 10 * time.Second
+	b.MaxElapsedTime = 0
+	return backoff.NewTickerWithTimer(backoff.WithContext(b, ctx), timer)
+}
+
 // Await signals the goroutine running Run that no more metrics will be sent on the channel.
 // Then it waits for that goroutine to submit them (with retries).
 func (s *Submitter) Await() {

internal/metrics/submitter/submitter_test.go

@@ -0,0 +1,141 @@
+package submitter
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stackrox/image-prefetcher/internal/metrics/gen"
+
+	"github.com/neilotoole/slogt"
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+type fakeClient struct {
+	failures int
+	calls    int
+}
+
+func (f *fakeClient) Submit(ctx context.Context, _ ...grpc.CallOption) (gen.Metrics_SubmitClient, error) {
+	f.calls++
+	if f.failures >= f.calls {
+		return nil, fmt.Errorf("failing as requested, %d calls, %d faiulres", f.calls, f.failures)
+	}
+	return &fakeSubmitClient{}, nil
+}
+
+type fakeSubmitClient struct {
+}
+
+func (f *fakeSubmitClient) Send(result *gen.Result) error {
+	return nil
+}
+
+func (f *fakeSubmitClient) CloseAndRecv() (*gen.Empty, error) {
+	return nil, nil
+}
+
+func (f *fakeSubmitClient) Header() (metadata.MD, error) {
+	panic("unimplemented")
+}
+
+func (f *fakeSubmitClient) Trailer() metadata.MD {
+	panic("unimplemented")
+}
+
+func (f *fakeSubmitClient) CloseSend() error {
+	panic("unimplemented")
+}
+
+func (f *fakeSubmitClient) Context() context.Context {
+	panic("unimplemented")
+}
+
+func (f *fakeSubmitClient) SendMsg(m any) error {
+	panic("unimplemented")
+}
+
+func (f *fakeSubmitClient) RecvMsg(m any) error {
+	panic("unimplemented")
+}
+
+type testTimer struct {
+	c chan time.Time
+}
+
+func (t *testTimer) Start(duration time.Duration) {
+	go func() { t.c <- time.Now().Add(duration) }()
+}
+
+func (t *testTimer) Stop() {
+}
+
+func (t *testTimer) C() <-chan time.Time {
+	return t.c
+}
+
+func TestSubmitter(t *testing.T) {
+	tests := map[string]struct {
+		client      *fakeClient
+		expectCalls int
+		timer       *testTimer
+		timeout     time.Duration
+		expectErr   error
+	}{
+		"nil": {
+			client: nil,
+		},
+		"simple": {
+			client:      &fakeClient{},
+			expectCalls: 1,
+		},
+		"with retries": {
+			client: &fakeClient{
+				failures: 2,
+			},
+			timer:       &testTimer{},
+			expectCalls: 3,
+		},
+		"timeout": {
+			client: &fakeClient{
+				failures: 999,
+			},
+			timeout:     50 * time.Millisecond,
+			expectCalls: 1,
+			expectErr:   context.DeadlineExceeded,
+		},
+	}
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			var sink *Submitter
+			if test.client != nil {
+				sink = NewSubmitter(slogt.New(t), test.client)
+				if test.timer != nil {
+					sink.timer = test.timer
+				}
+			}
+			timeout := test.timeout
+			if timeout == 0 {
+				timeout = 30 * time.Second // to catch hangs early
+			}
+			ctx, cancel := context.WithTimeout(context.Background(), timeout)
+			func() {
+				c := sink.Chan()
+				if sink != nil {
+					go func() { assert.ErrorIs(t, sink.Run(ctx), test.expectErr) }()
+					c <- &gen.Result{Error: "bam"}
+				}
+				sink.Await()
+				var actualCalls int
+				if test.client != nil {
+					actualCalls = test.client.calls
+				}
+				assert.Equal(t, test.expectCalls, actualCalls)
+				defer cancel()
+			}()
+		})
+	}
+}