From 1646931d34166c5404d19e180cffa64ed21d1e22 Mon Sep 17 00:00:00 2001
From: hatz <hatz@hey.com>
Date: Tue, 26 Nov 2024 12:45:21 -0600
Subject: [PATCH] Add a few cache related tests

---
 pkg/registry/ecr_test.go | 57 ++++++++++++++++++++++++++++++++++++++
 pkg/registry/gar_test.go | 59 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 115 insertions(+), 1 deletion(-)

diff --git a/pkg/registry/ecr_test.go b/pkg/registry/ecr_test.go
index 4660b96f..a6674d91 100644
--- a/pkg/registry/ecr_test.go
+++ b/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")
+			}
+		})
+	}
+}
diff --git a/pkg/registry/gar_test.go b/pkg/registry/gar_test.go
index 17ddde7e..f6ebd55a 100644
--- a/pkg/registry/gar_test.go
+++ b/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")
+			}
+		})
+	}
+}