Assignment 4: add unit, integration, and e2e tests; refactor for testability

internal/store/main_test.go

@@ -0,0 +1,61 @@
+package store
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"strings"
+	"testing"
+
+	"5803025015/w04/internal/database"
+
+	"github.com/jackc/pgx/v5"
+)
+
+// testConn is the shared connection for non-parallel store tests.
+var testConn *pgx.Conn
+
+func TestMain(m *testing.M) {
+	dbName := os.Getenv("DB_NAME")
+	if !strings.HasSuffix(dbName, "_test") {
+		fmt.Fprintf(os.Stderr, "REFUSING TO RUN: DB_NAME=%q must end with _test\n", dbName)
+		os.Exit(1)
+	}
+
+	conn, err := database.Connect()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "cannot connect to test DB: %v\n", err)
+		os.Exit(1)
+	}
+	testConn = conn
+
+	if err := database.CreateTables(testConn); err != nil {
+		fmt.Fprintf(os.Stderr, "cannot create tables: %v\n", err)
+		os.Exit(1)
+	}
+
+	code := m.Run()
+	_ = testConn.Close(context.Background())
+	os.Exit(code)
+}
+
+// truncateAll wipes every table. Call at the START of every test.
+func truncateAll(t *testing.T) {
+	t.Helper()
+	ctx := context.Background()
+	_, err := testConn.Exec(ctx, `TRUNCATE TABLE snippets, sessions, users CASCADE`)
+	if err != nil {
+		t.Fatalf("truncate failed: %v", err)
+	}
+}
+
+// newFreshConn opens a brand-new connection to the same DB (for restart tests).
+func newFreshConn(t *testing.T) *pgx.Conn {
+	t.Helper()
+	conn, err := database.Connect()
+	if err != nil {
+		t.Fatalf("cannot open fresh connection: %v", err)
+	}
+	t.Cleanup(func() { _ = conn.Close(context.Background()) })
+	return conn
+}