yellowjacket/backend/explore/artifactfetch_test.go

package explore

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/klauspost/compress/zstd"

	"yellowjacket/backend/database"
)

// artifactServer serves a compressed artifact and its checksum the way
// the Gitea generic package registry does.
func artifactServer(t *testing.T, body []byte) *httptest.Server {
	t.Helper()

	sum := sha256.Sum256(body)
	checksum := hex.EncodeToString(sum[:]) + "  " + coreArtifactFile + "\n"

	mux := http.NewServeMux()

	mux.HandleFunc("/"+coreArtifactChecksum,
		func(w http.ResponseWriter, _ *http.Request) {
			_, _ = w.Write([]byte(checksum))
		})

	// ServeContent gives the stub real Range support, so the resume path
	// is exercised against the same semantics as the package registry.
	mux.HandleFunc("/"+coreArtifactFile,
		func(w http.ResponseWriter, r *http.Request) {
			http.ServeContent(
				w, r, coreArtifactFile, time.Time{}, bytes.NewReader(body))
		})

	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)

	return srv
}

// compressArtifact zstd-compresses a file the way CI publishes it.
func compressArtifact(t *testing.T, path string) []byte {
	t.Helper()

	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("read artifact: %v", err)
	}

	enc, err := zstd.NewWriter(nil)
	if err != nil {
		t.Fatalf("zstd writer: %v", err)
	}

	defer func() { _ = enc.Close() }()

	return enc.EncodeAll(raw, nil)
}

// withArtifactEnv points the fetcher at a stub server and gives it an
// isolated data directory.
func withArtifactEnv(t *testing.T, baseURL string) {
	t.Helper()

	t.Setenv(artifactURLEnv, baseURL)
	t.Setenv("YJ_HOME", t.TempDir())
}

func TestFetchAndMergeArtifactEndToEnd(t *testing.T) {
	src := writeTestArtifact(t, validMeta(), []artifactRow{
		{"artist", artA, "Artist A", "Artist A", artA, 5000},
		{"release_group", rgA, "Album A", "Artist A", artA, 3000},
		{"recording", recA, "Song A", "Artist A", artA, 2000},
	})

	srv := artifactServer(t, compressArtifact(t, src))
	withArtifactEnv(t, srv.URL+"/")

	db := database.NewTestDB(t)
	si := NewSearchIndex(db, nil, nil, testLogger())

	if err := si.tryCoreArtifact(context.Background()); err != nil {
		t.Fatalf("tryCoreArtifact: %v", err)
	}

	var got int
	if err := db.QueryRowWriter(
		"SELECT COUNT(*) FROM explore_index",
	).Scan(&got); err != nil {
		t.Fatalf("count rows: %v", err)
	}

	if got != 3 {
		t.Errorf("merged %d rows, want 3", got)
	}

	if !si.artifactAlreadyMerged() {
		t.Error("artifact merge not recorded; a restart would re-download it")
	}

	// Staging must not keep a few hundred MB around after success.
	staging := filepath.Join(os.Getenv("YJ_HOME"), "data", "explore-staging")
	for _, name := range []string{coreArtifactFile, "core-index.db"} {
		if _, err := os.Stat(filepath.Join(staging, name)); err == nil {
			t.Errorf("%s left behind in staging after import", name)
		}
	}
}

func TestFetchArtifactRejectsBadChecksum(t *testing.T) {
	src := writeTestArtifact(t, validMeta(), []artifactRow{
		{"artist", artA, "Artist A", "Artist A", artA, 5000},
	})

	body := compressArtifact(t, src)

	// Serve a checksum for different content.
	mux := http.NewServeMux()
	mux.HandleFunc("/"+coreArtifactChecksum,
		func(w http.ResponseWriter, _ *http.Request) {
			_, _ = w.Write([]byte(strings.Repeat("a", 64) + "  " + coreArtifactFile))
		})
	mux.HandleFunc("/"+coreArtifactFile,
		func(w http.ResponseWriter, _ *http.Request) {
			_, _ = w.Write(body)
		})

	srv := httptest.NewServer(mux)
	defer srv.Close()

	withArtifactEnv(t, srv.URL+"/")

	db := database.NewTestDB(t)
	si := NewSearchIndex(db, nil, nil, testLogger())

	err := si.tryCoreArtifact(context.Background())
	if err == nil {
		t.Fatal("expected checksum rejection")
	}

	if !strings.Contains(err.Error(), "checksum mismatch") {
		t.Errorf("error = %v, want a checksum mismatch", err)
	}

	// A corrupt download must not leave the index claiming a catalog.
	if si.hasMeta(dumpImportDoneKey) || si.artifactAlreadyMerged() {
		t.Error("failed download still marked the catalog as imported")
	}
}

// A missing artifact is an ordinary outcome — the app has to keep
// working before the first one is ever published.
func TestFetchArtifactMissingIsUnavailable(t *testing.T) {
	srv := httptest.NewServer(http.NotFoundHandler())
	defer srv.Close()

	withArtifactEnv(t, srv.URL+"/")

	db := database.NewTestDB(t)
	si := NewSearchIndex(db, nil, nil, testLogger())

	err := si.tryCoreArtifact(context.Background())
	if !errors.Is(err, ErrArtifactUnavailable) {
		t.Errorf("error = %v, want ErrArtifactUnavailable", err)
	}
}

// The download resumes rather than restarting, which is what makes a
// large artifact survive a flaky connection.
func TestFetchArtifactResumesPartialDownload(t *testing.T) {
	src := writeTestArtifact(t, validMeta(), []artifactRow{
		{"artist", artA, "Artist A", "Artist A", artA, 5000},
	})

	body := compressArtifact(t, src)
	srv := artifactServer(t, body)
	withArtifactEnv(t, srv.URL+"/")

	db := database.NewTestDB(t)
	si := NewSearchIndex(db, nil, nil, testLogger())

	fetcher, err := newArtifactFetcher(si)
	if err != nil {
		t.Fatalf("newArtifactFetcher: %v", err)
	}

	// Pre-seed a truncated download.
	half := len(body) / 2
	if err := os.WriteFile(fetcher.compressedPath(), body[:half], 0o644); err != nil {
		t.Fatalf("seed partial download: %v", err)
	}

	if _, err := fetcher.fetch(context.Background()); err != nil {
		t.Fatalf("fetch after partial: %v", err)
	}
}