Add a few cache related tests

pkg/registry/ecr_test.go

@@ -1,10 +1,13 @@
 package registry
 
 import (
+	"context"
 	"encoding/base64"
 	"testing"
+	"time"
 
 	"github.com/containers/image/v5/transports/alltransports"
+	"github.com/dgraph-io/ristretto"
 
 	"github.com/estahn/k8s-image-swapper/pkg/config"
 	"github.com/stretchr/testify/assert"
@@ -51,3 +54,57 @@ func TestECRIsOrigin(t *testing.T) {
 		assert.Equal(t, testcase.expected, result)
 	}
 }
+
+func TestEcrImageExistsCaching(t *testing.T) {
+	// Setup a test cache
+	cache, err := ristretto.NewCache(&ristretto.Config{
+		NumCounters: 1e7,     // number of keys to track frequency of (10M).
+		MaxCost:     1 << 30, // maximum cost of cache (1GB).
+		BufferItems: 64,      // number of keys per Get buffer.
+	})
+	assert.NoError(t, err)
+
+	tests := []struct {
+		name            string
+		cacheTtlMinutes int
+		expectCached    bool
+	}{
+		{
+			name:            "cache disabled when TTL is 0",
+			cacheTtlMinutes: 0,
+			expectCached:    false,
+		},
+		{
+			name:            "cache enabled with TTL and jitter",
+			cacheTtlMinutes: 60,
+			expectCached:    true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			client := NewDummyECRClient("us-east-1", "12345678912", "", config.ECROptions{}, []byte(""))
+
+			// Setup cache
+			client.cache = cache
+			client.cacheTtlMinutes = tc.cacheTtlMinutes
+
+			// Create a test image reference and add to cache. Use 100ms as TTL
+			imageRef, err := alltransports.ParseImageName("docker://12345678912.dkr.ecr.us-east-1.amazonaws.com/test-project/repo/test-image:latest")
+			cache.SetWithTTL(imageRef.DockerReference().String(), true, 1, 100*time.Millisecond)
+			assert.NoError(t, err)
+
+			// Cache should be a hit
+			exists := client.ImageExists(ctx, imageRef)
+			assert.Equal(t, tc.expectCached, exists)
+
+			if tc.expectCached {
+				// Verify cache expiry
+				time.Sleep(time.Duration(150 * time.Millisecond)) // Use milliseconds for testing
+				_, found := client.cache.Get(imageRef.DockerReference().String())
+				assert.False(t, found, "cache entry should have expired")
+			}
+		})
+	}
+}

pkg/registry/gar_test.go

@@ -1,10 +1,12 @@
 package registry
 
 import (
+	"context"
 	"testing"
+	"time"
 
 	"github.com/containers/image/v5/transports/alltransports"
-
+	"github.com/dgraph-io/ristretto"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -36,3 +38,58 @@ func TestGARIsOrigin(t *testing.T) {
 		assert.Equal(t, testcase.expected, result)
 	}
 }
+
+func TestGarImageExistsCaching(t *testing.T) {
+	// Setup a test cache
+	cache, err := ristretto.NewCache(&ristretto.Config{
+		NumCounters: 1e7,     // number of keys to track frequency of (10M).
+		MaxCost:     1 << 30, // maximum cost of cache (1GB).
+		BufferItems: 64,      // number of keys per Get buffer.
+	})
+	assert.NoError(t, err)
+
+	tests := []struct {
+		name            string
+		cacheTtlMinutes int
+		expectCached    bool
+	}{
+		{
+			name:            "cache disabled when TTL is 0",
+			cacheTtlMinutes: 0,
+			expectCached:    false,
+		},
+		{
+			name:            "cache enabled with TTL and jitter",
+			cacheTtlMinutes: 60,
+			expectCached:    true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			client, err := NewMockGARClient(nil, "us-central1-docker.pkg.dev/test-project/repo")
+			assert.NoError(t, err)
+
+			// Setup cache
+			client.cache = cache
+			client.cacheTtlMinutes = tc.cacheTtlMinutes
+
+			// Create a test image reference and add to cache. Use 100ms as TTL
+			imageRef, err := alltransports.ParseImageName("docker://us-central1-docker.pkg.dev/test-project/repo/test-image:latest")
+			cache.SetWithTTL(imageRef.DockerReference().String(), true, 1, 100*time.Millisecond)
+			assert.NoError(t, err)
+
+			// Cache should be a hit
+			exists := client.ImageExists(ctx, imageRef)
+			assert.Equal(t, tc.expectCached, exists)
+
+			if tc.expectCached {
+				// Verify cache expiry
+				time.Sleep(time.Duration(150 * time.Millisecond)) // Use milliseconds for testing
+				_, found := client.cache.Get(imageRef.DockerReference().String())
+				assert.False(t, found, "cache entry should have expired")
+			}
+		})
+	}
+}