pkg-proxy/internal/server/retention_pypi_test.go

package server

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

	"github.com/git-pkgs/proxy/internal/config"
	"github.com/git-pkgs/proxy/internal/database"
	"github.com/git-pkgs/proxy/internal/handler"
	"github.com/git-pkgs/proxy/internal/retention"
	"github.com/git-pkgs/registries/fetch"
)

const (
	// pypiProxyPrefix is the route the proxy rewrites download links to: it
	// keeps the upstream path, /packages/... included, after /packages.
	pypiProxyPrefix = "/packages"
	pypiWheelPath   = "/packages/ab/cd/0123/Typing_Extensions-4.0.0-py3-none-any.whl"
	pypiSdistPath   = "/packages/ef/01/4567/requests-2.31.0.tar.gz"
)

// pypiRetentionUpstream serves a wheel and an sdist, counting downloads.
func pypiRetentionUpstream(t *testing.T) (*httptest.Server, *atomic.Int64) {
	t.Helper()
	var downloads atomic.Int64
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case pypiWheelPath, pypiSdistPath:
			downloads.Add(1)
			_, _ = io.WriteString(w, "distribution bytes")
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(upstream.Close)
	return upstream, &downloads
}

func newPyPIRetentionProxy(t *testing.T, upstream *httptest.Server) (*httptest.Server, *database.DB) {
	t.Helper()
	db, store := setupEvictionTest(t)
	// The upstream's own client reaches the loopback test server, which the
	// default fetcher refuses.
	fetcher := fetch.NewFetcher(fetch.WithHTTPClient(upstream.Client()), fetch.WithMaxRetries(0))
	t.Cleanup(func() { _ = fetcher.Close() })
	proxy := handler.NewProxy(db, store, fetcher, fetch.NewResolver(), discardLogger())
	proxy.HTTPClient = upstream.Client()
	h := handler.NewPyPIHandlerWithUpstreams(proxy, "http://proxy.test", upstream.URL, upstream.URL)
	srv := httptest.NewServer(h.Routes())
	t.Cleanup(srv.Close)
	return srv, db
}

func pypiGet(t *testing.T, url string) {
	t.Helper()
	resp, err := http.Get(url)
	if err != nil {
		t.Fatalf("GET %s: %v", url, err)
	}
	body, _ := io.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("GET %s = %d %s, want 200", url, resp.StatusCode, body)
	}
}

// pypiArtifact returns the record of the file named filename.
func pypiArtifact(t *testing.T, db *database.DB, filename string) database.Artifact {
	t.Helper()
	var art database.Artifact
	err := db.Get(&art, db.Rebind(`SELECT * FROM artifacts WHERE filename = ?`), filename)
	if err != nil {
		t.Fatalf("reading artifact %s: %v", filename, err)
	}
	return art
}

func pypiBackdate(t *testing.T, db *database.DB, filename string, fetched, accessed time.Time) {
	t.Helper()
	query := db.Rebind(`UPDATE artifacts SET fetched_at = ?, last_accessed_at = ? WHERE filename = ?`)
	if _, err := db.Exec(query, fetched, accessed, filename); err != nil {
		t.Fatalf("backdating %s: %v", filename, err)
	}
}

func TestPyPIRegistersForRetention(t *testing.T) {
	if _, ok := retention.Default.Lookup("pypi"); !ok {
		t.Fatal("pypi is not registered for retention")
	}
	if key, ok := retention.Default.KeyForDBEcosystem("pypi"); !ok || key != "pypi" {
		t.Errorf("KeyForDBEcosystem(pypi) = %q, %v, want pypi, true", key, ok)
	}
}

// A PyPI file served from the cache stays while it is being used, is
// evicted once its rule runs out, comes back from upstream on the next
// download, and the fresh fetch keeps it even though the evicted copy's old
// access time is still on the record.
func TestPyPIRetentionEvictsAndRefetches(t *testing.T) {
	upstream, downloads := pypiRetentionUpstream(t)
	srv, db := newPyPIRetentionProxy(t, upstream)
	const filename = "requests-2.31.0.tar.gz"

	pypiGet(t, srv.URL+pypiProxyPrefix+pypiSdistPath)
	pypiGet(t, srv.URL+pypiProxyPrefix+pypiSdistPath)
	if got := downloads.Load(); got != 1 {
		t.Fatalf("upstream downloads = %d, want 1 (second from cache)", got)
	}
	if art := pypiArtifact(t, db, filename); !art.LastAccessedAt.Valid {
		t.Fatal("cache hit did not record an access")
	}

	rules := retention.Rules{Ecosystems: map[string]time.Duration{"pypi": retentionDay}}
	now := time.Now()
	sweepRetention(context.Background(), db, discardLogger(), retention.Default, rules, now)
	if !pypiArtifact(t, db, filename).StoragePath.Valid {
		t.Fatal("file used today was evicted")
	}

	sweepRetention(context.Background(), db, discardLogger(), retention.Default, rules, now.Add(2*retentionDay))
	if pypiArtifact(t, db, filename).StoragePath.Valid {
		t.Fatal("file unused for two days kept under a one-day rule")
	}

	pypiGet(t, srv.URL+pypiProxyPrefix+pypiSdistPath)
	if got := downloads.Load(); got != 2 {
		t.Errorf("upstream downloads = %d, want 2 after eviction", got)
	}
	old := time.Now().Add(-30 * retentionDay)
	art := pypiArtifact(t, db, filename)
	pypiBackdate(t, db, filename, art.FetchedAt.Time, old)
	sweepRetention(context.Background(), db, discardLogger(), retention.Default, rules, time.Now())
	if !pypiArtifact(t, db, filename).StoragePath.Valid {
		t.Error("refetched file evicted again by the next sweep")
	}
}

// Wheel names use underscores and the publisher's casing; a package rule
// written in any of the spellings PyPI treats as the same project has to
// reach the stored package.
func TestPyPIRetentionPackageRuleMatchesStoredPackage(t *testing.T) {
	upstream, _ := pypiRetentionUpstream(t)
	srv, db := newPyPIRetentionProxy(t, upstream)
	pypiGet(t, srv.URL+pypiProxyPrefix+pypiWheelPath)
	pypiGet(t, srv.URL+pypiProxyPrefix+pypiSdistPath)

	wheel := pypiArtifact(t, db, "Typing_Extensions-4.0.0-py3-none-any.whl")
	var stored string
	if err := db.Get(&stored, db.Rebind(`SELECT package_purl FROM versions WHERE purl = ?`), wheel.VersionPURL); err != nil {
		t.Fatalf("reading stored package PURL: %v", err)
	}

	for _, key := range []string{"pkg:pypi/typing-extensions", "pkg:pypi/Typing_Extensions", "pkg:pypi/typing_extensions"} {
		cfg := config.Default()
		cfg.Storage.Retention.Packages = map[string]string{key: "0"}
		rules, err := cfg.RetentionRules(retention.Default)
		if err != nil {
			t.Fatalf("RetentionRules(%s): %v", key, err)
		}
		if _, ok := rules.Packages[stored]; !ok {
			t.Errorf("rule %s resolved to %v, want the stored package %s", key, rules.Packages, stored)
		}
	}

	// With both files long unused, the "0" rule keeps the wheel while the
	// ecosystem rule evicts the requests sdist.
	old := time.Now().Add(-30 * retentionDay)
	for _, filename := range []string{"Typing_Extensions-4.0.0-py3-none-any.whl", "requests-2.31.0.tar.gz"} {
		pypiBackdate(t, db, filename, old, old)
	}
	cfg := config.Default()
	cfg.Storage.Retention.Ecosystems = map[string]string{"pypi": "1d"}
	cfg.Storage.Retention.Packages = map[string]string{"pkg:pypi/Typing_Extensions": "0"}
	rules, err := cfg.RetentionRules(retention.Default)
	if err != nil {
		t.Fatalf("RetentionRules: %v", err)
	}
	sweepRetention(context.Background(), db, discardLogger(), retention.Default, rules, time.Now())

	if !pypiArtifact(t, db, "Typing_Extensions-4.0.0-py3-none-any.whl").StoragePath.Valid {
		t.Error("typing-extensions evicted despite its never-evict rule")
	}
	if pypiArtifact(t, db, "requests-2.31.0.tar.gz").StoragePath.Valid {
		t.Error("requests kept past the pypi rule")
	}
	if !strings.HasPrefix(stored, "pkg:pypi/") {
		t.Errorf("stored package PURL = %q, want a pkg:pypi PURL", stored)
	}
}