tiny-proxy/internal/proxycache/proxycache_test.go

package proxycache

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"tiny-proxy/internal/metrics"
	"tiny-proxy/internal/storage"
)

func testCache(t *testing.T) *Cache {
	t.Helper()
	store, err := storage.NewLocalFS(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	return New(store)
}

func get(t *testing.T, c *Cache, upstream http.Handler, key string, ttl time.Duration) (int, string) {
	t.Helper()
	return getPolicy(t, c, upstream, key, ttl, false)
}

// getPolicy is get with an explicit NoStore policy on the fetch result.
func getPolicy(t *testing.T, c *Cache, upstream http.Handler, key string, ttl time.Duration, noStore bool) (int, string) {
	t.Helper()
	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	c.Serve(rec, req, key, ttl, func(ctx context.Context) (Result, error) {
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "", nil)
		rec2 := httptest.NewRecorder()
		upstream.ServeHTTP(rec2, req)
		return Result{Resp: rec2.Result(), NoStore: noStore}, nil
	})
	res := rec.Result()
	body, _ := io.ReadAll(res.Body)
	return res.StatusCode, string(body)
}

func TestMissThenHit(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		w.Header().Set("Content-Type", "text/plain")
		fmt.Fprint(w, "payload")
	})
	c := testCache(t)

	code, body := get(t, c, upstream, "k", 0)
	if code != 200 || body != "payload" {
		t.Fatalf("miss: code=%d body=%q", code, body)
	}
	code, body = get(t, c, upstream, "k", 0)
	if code != 200 || body != "payload" {
		t.Fatalf("hit: code=%d body=%q", code, body)
	}
	if hits.Load() != 1 {
		t.Errorf("upstream hits = %d, want 1 (second request should be cached)", hits.Load())
	}
}

func TestTTLRefetch(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "v%d", hits.Add(1))
	})
	c := testCache(t)

	// Seed the cache with an object stored an hour ago.
	if err := c.Store.Put(t.Context(), "k", strings.NewReader("v0"),
		storage.ObjectMeta{StoredAt: time.Now().Add(-time.Hour)}); err != nil {
		t.Fatal(err)
	}
	_, body := get(t, c, upstream, "k", time.Minute)
	if body != "v1" {
		t.Errorf("expired entry should refetch, body = %q", body)
	}
	// The refetch refreshed the entry: next request is served from cache.
	_, body = get(t, c, upstream, "k", time.Minute)
	if body != "v1" || hits.Load() != 1 {
		t.Errorf("refreshed entry should hit cache, body = %q hits = %d", body, hits.Load())
	}
}

// TestNoStoreServesWithoutCaching covers upstream cache: false — the body
// reaches the client on every request but never lands in storage.
func TestNoStoreServesWithoutCaching(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		fmt.Fprint(w, "payload")
	})
	c := testCache(t)

	for i := 0; i < 2; i++ {
		code, body := getPolicy(t, c, upstream, "k", 0, true)
		if code != http.StatusOK || body != "payload" {
			t.Fatalf("request %d: code = %d body = %q, want 200 %q", i, code, body, "payload")
		}
	}
	if got := hits.Load(); got != 2 {
		t.Errorf("upstream hits = %d, want 2 (nothing may be cached)", got)
	}
	if _, err := c.Store.Stat(t.Context(), "k"); err != storage.ErrNotFound {
		t.Error("no-store response must not be written to storage")
	}
}

// TestNoStoreStillServesCachedCopy: the flag suppresses writes, not reads, so
// an object cached earlier (e.g. by another upstream of the same vrepo) is
// still served without touching the upstream.
func TestNoStoreStillServesCachedCopy(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		fmt.Fprint(w, "payload")
	})
	c := testCache(t)

	get(t, c, upstream, "k", 0) // cached by a normal (cacheable) upstream
	if _, body := getPolicy(t, c, upstream, "k", 0, true); body != "payload" {
		t.Errorf("body = %q, want %q", body, "payload")
	}
	if got := hits.Load(); got != 1 {
		t.Errorf("upstream hits = %d, want 1 (second request served from cache)", got)
	}
}

func TestFetchNoStore(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		fmt.Fprint(w, "value")
	})
	c := testCache(t)

	for i := 0; i < 2; i++ {
		body, err := fetchBodyPolicy(c, "k", time.Hour, upstream, true)
		if err != nil || string(body) != "value" {
			t.Fatalf("request %d: body = %q err = %v", i, body, err)
		}
	}
	if got := hits.Load(); got != 2 {
		t.Errorf("upstream hits = %d, want 2", got)
	}
	if _, err := c.Store.Stat(context.Background(), "k"); err != storage.ErrNotFound {
		t.Error("no-store fetch must not be written to storage")
	}
}

func TestUpstreamErrorNotCached(t *testing.T) {
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "gone", http.StatusNotFound)
	})
	c := testCache(t)
	code, _ := get(t, c, upstream, "k", 0)
	if code != http.StatusNotFound {
		t.Errorf("code = %d, want 404 passthrough", code)
	}
	if _, err := c.Store.Stat(t.Context(), "k"); err != storage.ErrNotFound {
		t.Error("404 response should not be cached")
	}
}

func fetchBody(c *Cache, key string, ttl time.Duration, upstream http.Handler) ([]byte, error) {
	return fetchBodyPolicy(c, key, ttl, upstream, false)
}

func fetchBodyPolicy(c *Cache, key string, ttl time.Duration, upstream http.Handler, noStore bool) ([]byte, error) {
	return c.Fetch(context.Background(), key, ttl, func(ctx context.Context) (Result, error) {
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "", nil)
		rec := httptest.NewRecorder()
		upstream.ServeHTTP(rec, req)
		return Result{Resp: rec.Result(), NoStore: noStore}, nil
	})
}

func TestFetchMissThenHit(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		fmt.Fprint(w, "value")
	})
	c := testCache(t)

	body, err := fetchBody(c, "k", 0, upstream)
	if err != nil || string(body) != "value" {
		t.Fatalf("miss: body=%q err=%v", body, err)
	}
	body, err = fetchBody(c, "k", 0, upstream)
	if err != nil || string(body) != "value" {
		t.Fatalf("hit: body=%q err=%v", body, err)
	}
	if hits.Load() != 1 {
		t.Errorf("upstream hits = %d, want 1", hits.Load())
	}
}

func TestFetchTTLRefetch(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "v%d", hits.Add(1))
	})
	c := testCache(t)
	if err := c.Store.Put(t.Context(), "k", strings.NewReader("v0"),
		storage.ObjectMeta{StoredAt: time.Now().Add(-time.Hour)}); err != nil {
		t.Fatal(err)
	}
	body, err := fetchBody(c, "k", time.Minute, upstream)
	if err != nil || string(body) != "v1" {
		t.Errorf("expired entry should refetch, body=%q err=%v", body, err)
	}
	body, _ = fetchBody(c, "k", time.Minute, upstream)
	if string(body) != "v1" || hits.Load() != 1 {
		t.Errorf("refreshed entry should hit cache, body=%q hits=%d", body, hits.Load())
	}
}

func TestFetchNon200NotCached(t *testing.T) {
	var hits atomic.Int32
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		http.Error(w, "not here", http.StatusNotFound)
	})
	c := testCache(t)
	_, err := fetchBody(c, "k", 0, upstream)
	var ue *UpstreamError
	if !errors.As(err, &ue) || ue.StatusCode != http.StatusNotFound {
		t.Fatalf("err = %v, want *UpstreamError with 404", err)
	}
	if _, serr := c.Store.Stat(t.Context(), "k"); serr != storage.ErrNotFound {
		t.Error("non-200 response should not be cached")
	}
	// Next Fetch retries upstream.
	if _, err := fetchBody(c, "k", 0, upstream); err == nil || hits.Load() != 2 {
		t.Errorf("hits = %d, want 2 (no caching of failures)", hits.Load())
	}
}

// failingStore fails every Put immediately, without reading the body.
type failingStore struct{ storage.Store }

func (failingStore) Put(context.Context, string, io.Reader, storage.ObjectMeta) error {
	return errors.New("storage is read-only")
}

func TestMetricsMissThenHit(t *testing.T) {
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "payload")
	})
	c := testCache(t)
	c.Metrics = metrics.NewRegistry()

	get(t, c, upstream, "k", 0)
	get(t, c, upstream, "k", 0)
	hits, misses, errors := c.Metrics.Snapshot()
	if hits != 1 || misses != 1 || errors != 0 {
		t.Errorf("snapshot = hits:%d misses:%d errors:%d, want 1/1/0", hits, misses, errors)
	}
}

func TestMetricsFetchError(t *testing.T) {
	c := testCache(t)
	c.Metrics = metrics.NewRegistry()
	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	c.Serve(rec, req, "k", 0, func(ctx context.Context) (Result, error) {
		return Result{}, errors.New("upstream down")
	})
	_, misses, errors := c.Metrics.Snapshot()
	if misses != 1 || errors != 1 {
		t.Errorf("snapshot misses:%d errors:%d, want 1/1", misses, errors)
	}
}

func TestMetricsStoreError(t *testing.T) {
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "payload")
	})
	c := New(failingStore{testCache(t).Store})
	c.Metrics = metrics.NewRegistry()
	get(t, c, upstream, "k", 0)
	_, misses, errors := c.Metrics.Snapshot()
	if misses != 1 || errors != 1 {
		t.Errorf("snapshot misses:%d errors:%d, want 1/1", misses, errors)
	}
}

func TestMetricsFetchPath(t *testing.T) {
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "value")
	})
	c := testCache(t)
	c.Metrics = metrics.NewRegistry()
	fetchBody(c, "k", 0, upstream)
	fetchBody(c, "k", 0, upstream)
	hits, misses, errors := c.Metrics.Snapshot()
	if hits != 1 || misses != 1 || errors != 0 {
		t.Errorf("snapshot = hits:%d misses:%d errors:%d, want 1/1/0", hits, misses, errors)
	}
}

func TestPutFailureDoesNotHang(t *testing.T) {
	payload := strings.Repeat("x", 1<<20) // larger than any pipe buffer
	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, payload)
	})
	c := New(failingStore{testCache(t).Store})

	done := make(chan string, 1)
	go func() {
		_, body := get(t, c, upstream, "k", 0)
		done <- body
	}()
	select {
	case body := <-done:
		if body != payload {
			t.Errorf("body length = %d, want %d (client should still get the full body)", len(body), len(payload))
		}
	case <-time.After(10 * time.Second):
		t.Fatal("request hung: Put failing before reading the pipe must not deadlock the client")
	}
}