ncro/internal/router/router_test.go

package router_test

import (
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"

	"notashelf.dev/ncro/internal/cache"
	"notashelf.dev/ncro/internal/prober"
	"notashelf.dev/ncro/internal/router"
)

func newTestRouter(t *testing.T, upstreams ...string) (*router.Router, func()) {
	t.Helper()
	f, _ := os.CreateTemp("", "ncro-router-*.db")
	f.Close()
	db, err := cache.Open(f.Name(), 1000)
	if err != nil {
		t.Fatal(err)
	}
	p := prober.New(0.3)
	for _, u := range upstreams {
		p.RecordLatency(u, 10)
	}
	r := router.New(db, p, time.Hour, 5*time.Second)
	return r, func() {
		db.Close()
		os.Remove(f.Name())
	}
}

func TestRouteHit(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "StorePath: /nix/store/abc123-hello")
	}))
	defer srv.Close()

	r, cleanup := newTestRouter(t, srv.URL)
	defer cleanup()

	result, err := r.Resolve("abc123", []string{srv.URL})
	if err != nil {
		t.Fatalf("Resolve: %v", err)
	}
	if result.URL != srv.URL {
		t.Errorf("url = %q, want %q", result.URL, srv.URL)
	}
	if result.LatencyMs <= 0 {
		t.Error("expected positive latency")
	}
}

func TestRouteRacePicksFastest(t *testing.T) {
	fast := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
	}))
	defer fast.Close()

	slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(100 * time.Millisecond)
		w.WriteHeader(200)
	}))
	defer slow.Close()

	r, cleanup := newTestRouter(t, fast.URL, slow.URL)
	defer cleanup()

	result, err := r.Resolve("somehash", []string{slow.URL, fast.URL})
	if err != nil {
		t.Fatalf("Resolve: %v", err)
	}
	if result.URL != fast.URL {
		t.Errorf("expected fast server to win, got %q", result.URL)
	}
}

func TestRouteAllFail(t *testing.T) {
	r, cleanup := newTestRouter(t)
	defer cleanup()

	_, err := r.Resolve("somehash", []string{"http://127.0.0.1:1"})
	if err == nil {
		t.Error("expected error when all upstreams fail")
	}
}

func TestRouteAllNotFound(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(404)
	}))
	defer srv.Close()

	r, cleanup := newTestRouter(t, srv.URL)
	defer cleanup()

	_, err := r.Resolve("somehash", []string{srv.URL})
	if !errors.Is(err, router.ErrNotFound) {
		t.Errorf("expected ErrNotFound, got %v", err)
	}
}

func TestRouteAllUnavailable(t *testing.T) {
	r, cleanup := newTestRouter(t)
	defer cleanup()

	_, err := r.Resolve("somehash", []string{"http://127.0.0.1:1"})
	if !errors.Is(err, router.ErrUpstreamUnavailable) {
		t.Errorf("expected ErrUpstreamUnavailable, got %v", err)
	}
}

func TestRaceWithMalformedURL(t *testing.T) {
	r, cleanup := newTestRouter(t)
	defer cleanup()

	_, err := r.Resolve("somehash", []string{"://bad-url"})
	if err == nil {
		t.Error("expected error for malformed upstream URL")
	}
}

func TestCacheHit(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
	}))
	defer srv.Close()

	r, cleanup := newTestRouter(t, srv.URL)
	defer cleanup()

	r.Resolve("abc123", []string{srv.URL})

	result, err := r.Resolve("abc123", []string{srv.URL})
	if err != nil {
		t.Fatalf("second Resolve: %v", err)
	}
	if !result.CacheHit {
		t.Error("expected cache hit on second resolve")
	}
}

func TestResolveWithDownUpstream(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
	}))
	defer srv.Close()

	f, _ := os.CreateTemp("", "ncro-router-*.db")
	f.Close()
	db, _ := cache.Open(f.Name(), 1000)
	defer db.Close()
	defer os.Remove(f.Name())

	p := prober.New(0.3)
	p.RecordLatency(srv.URL, 10)
	// Force the upstream to StatusDown
	for range 10 {
		p.RecordFailure(srv.URL)
	}

	r := router.New(db, p, time.Hour, 5*time.Second)
	// Router should still attempt the race (the race uses HEAD, not the prober status)
	// The upstream is actually healthy (httptest), so the race should succeed.
	result, err := r.Resolve("somehash", []string{srv.URL})
	if err != nil {
		t.Fatalf("Resolve with down-flagged upstream: %v", err)
	}
	if result.URL != srv.URL {
		t.Errorf("url = %q", result.URL)
	}
}