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

2026-10-03 19:09:40 +07:00
parent 4ac21847fc
commit 3945920abb
24 changed files with 2753 additions and 147 deletions
@@ -2,10 +2,11 @@ package middleware

 import (
 	"context"
+	"encoding/json"
 	"net/http"
 	"strings"

-	"5803025015/w04/internal/store"
+	"5803025015/w04/internal/models"
 )

 type contextKey string
@@ -19,35 +20,46 @@ func GetUserID(ctx context.Context) string {
 	return ""
 }

-type AuthMiddleware struct {
-	sessionStore *store.SessionStore
+type SessionStore interface {
+	Get(ctx context.Context, token string) (*models.Session, error)
 }

-func NewAuthMiddleware(sessionStore *store.SessionStore) *AuthMiddleware {
+type AuthMiddleware struct {
+	sessionStore SessionStore
+}
+
+func NewAuthMiddleware(sessionStore SessionStore) *AuthMiddleware {
 	return &AuthMiddleware{sessionStore: sessionStore}
 }

+func writeJSONError(w http.ResponseWriter, status int, msg string) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(status)
+	_ = json.NewEncoder(w).Encode(map[string]string{"error": msg})
+}
+
 func (m *AuthMiddleware) RequireAuth(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		authHeader := r.Header.Get("Authorization")
 		if authHeader == "" {
-			http.Error(w, "Unauthorized", http.StatusUnauthorized)
+			writeJSONError(w, http.StatusUnauthorized, "unauthorized")
 			return
 		}

-		token := authHeader
-		if strings.HasPrefix(strings.ToLower(authHeader), "bearer ") {
-			token = authHeader[7:]
+		const bearerPrefix = "bearer "
+		if !strings.HasPrefix(strings.ToLower(authHeader), bearerPrefix) {
+			writeJSONError(w, http.StatusUnauthorized, "unauthorized")
+			return
 		}
-		token = strings.TrimSpace(token)
+		token := strings.TrimSpace(authHeader[len(bearerPrefix):])
 		if token == "" {
-			http.Error(w, "Unauthorized", http.StatusUnauthorized)
+			writeJSONError(w, http.StatusUnauthorized, "unauthorized")
 			return
 		}

 		session, err := m.sessionStore.Get(r.Context(), token)
 		if err != nil || session == nil {
-			http.Error(w, "Unauthorized", http.StatusUnauthorized)
+			writeJSONError(w, http.StatusUnauthorized, "unauthorized")
 			return
 		}

@@ -62,6 +74,6 @@ func (m *AuthMiddleware) RequireAuthFunc(next http.HandlerFunc) http.HandlerFunc
 	}
 }

-func RequireAuth(sessionStore *store.SessionStore) func(http.Handler) http.Handler {
+func RequireAuth(sessionStore SessionStore) func(http.Handler) http.Handler {
 	return NewAuthMiddleware(sessionStore).RequireAuth
 }
@@ -0,0 +1,217 @@
+package middleware
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"5803025015/w04/internal/models"
+)
+
+// ---------- hand-written fake ----------
+
+// fakeSessionStore is a hand-written in-memory SessionStore for tests.
+type fakeSessionStore struct {
+	sessions map[string]*models.Session
+	err      error // if set, Get returns this error
+}
+
+func newFakeSessionStore() *fakeSessionStore {
+	return &fakeSessionStore{sessions: make(map[string]*models.Session)}
+}
+
+func (f *fakeSessionStore) Get(ctx context.Context, token string) (*models.Session, error) {
+	if f.err != nil {
+		return nil, f.err
+	}
+	s, ok := f.sessions[token]
+	if !ok {
+		return nil, errors.New("not found")
+	}
+	return s, nil
+}
+
+// ---------- helpers ----------
+
+// callMiddleware runs RequireAuth around a next handler that records whether
+// it was called and what user id it saw.
+func callMiddleware(
+	t *testing.T,
+	store SessionStore,
+	authHeader string,
+) (rec *httptest.ResponseRecorder, called *bool, seenUserID *string) {
+	t.Helper()
+
+	calledVal := false
+	seenID := ""
+
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		calledVal = true
+		seenID = GetUserID(r.Context())
+		w.WriteHeader(http.StatusOK)
+	})
+
+	mw := NewAuthMiddleware(store)
+	handler := mw.RequireAuth(next)
+
+	req := httptest.NewRequest(http.MethodGet, "/protected", nil)
+	if authHeader != "" {
+		req.Header.Set("Authorization", authHeader)
+	}
+	rec = httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	return rec, &calledVal, &seenID
+}
+
+// assertJSONError asserts that the body is valid JSON and has "error" key.
+func assertJSONError(t *testing.T, body []byte) {
+	t.Helper()
+	var m map[string]any
+	if err := json.Unmarshal(body, &m); err != nil {
+		t.Fatalf("body is not valid JSON: %v; body=%q", err, string(body))
+	}
+	if _, ok := m["error"]; !ok {
+		t.Fatalf("expected 'error' key in JSON body; got %v", m)
+	}
+}
+
+// ---------- tests ----------
+
+func TestRequireAuth(t *testing.T) {
+	t.Run("no Authorization header -> 401, next not called", func(t *testing.T) {
+		store := newFakeSessionStore()
+		rec, called, _ := callMiddleware(t, store, "")
+		if rec.Code != http.StatusUnauthorized {
+			t.Fatalf("want 401, got %d", rec.Code)
+		}
+		if *called {
+			t.Fatal("next handler should NOT have been called")
+		}
+		if ct := rec.Header().Get("Content-Type"); ct != "application/json" {
+			t.Fatalf("want Content-Type application/json, got %q", ct)
+		}
+		assertJSONError(t, rec.Body.Bytes())
+	})
+
+	t.Run("Basic scheme -> 401", func(t *testing.T) {
+		store := newFakeSessionStore()
+		rec, called, _ := callMiddleware(t, store, "Basic abc")
+		if rec.Code != http.StatusUnauthorized {
+			t.Fatalf("want 401, got %d", rec.Code)
+		}
+		if *called {
+			t.Fatal("next handler should NOT have been called")
+		}
+		assertJSONError(t, rec.Body.Bytes())
+		t.Run("RequireAuthFunc wrapper works", func(t *testing.T) {
+			store := newFakeSessionStore()
+			store.sessions["tok"] = &models.Session{Token: "tok", UserID: "u1"}
+
+			called := false
+			next := func(w http.ResponseWriter, r *http.Request) {
+				called = true
+				if GetUserID(r.Context()) != "u1" {
+					t.Errorf("wrong user id in context")
+				}
+				w.WriteHeader(http.StatusOK)
+			}
+
+			mw := NewAuthMiddleware(store)
+			h := mw.RequireAuthFunc(next)
+
+			req := httptest.NewRequest(http.MethodGet, "/x", nil)
+			req.Header.Set("Authorization", "Bearer tok")
+			rec := httptest.NewRecorder()
+			h.ServeHTTP(rec, req)
+
+			if rec.Code != http.StatusOK || !called {
+				t.Fatalf("want 200 + called, got %d called=%v", rec.Code, called)
+			}
+		})
+
+		t.Run("RequireAuth top-level function works", func(t *testing.T) {
+			store := newFakeSessionStore()
+			store.sessions["tok"] = &models.Session{Token: "tok", UserID: "u1"}
+
+			called := false
+			next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				called = true
+				w.WriteHeader(http.StatusOK)
+			})
+
+			mwFunc := RequireAuth(store)
+			h := mwFunc(next)
+
+			req := httptest.NewRequest(http.MethodGet, "/x", nil)
+			req.Header.Set("Authorization", "Bearer tok")
+			rec := httptest.NewRecorder()
+			h.ServeHTTP(rec, req)
+
+			if rec.Code != http.StatusOK || !called {
+				t.Fatalf("want 200 + called, got %d called=%v", rec.Code, called)
+			}
+		})
+	})
+
+	t.Run("Bearer with empty token -> 401", func(t *testing.T) {
+		store := newFakeSessionStore()
+		rec, called, _ := callMiddleware(t, store, "Bearer ")
+		if rec.Code != http.StatusUnauthorized {
+			t.Fatalf("want 401, got %d", rec.Code)
+		}
+		if *called {
+			t.Fatal("next handler should NOT have been called")
+		}
+		assertJSONError(t, rec.Body.Bytes())
+	})
+
+	t.Run("unknown token -> 401", func(t *testing.T) {
+		store := newFakeSessionStore()
+		rec, called, _ := callMiddleware(t, store, "Bearer not-a-real-token")
+		if rec.Code != http.StatusUnauthorized {
+			t.Fatalf("want 401, got %d", rec.Code)
+		}
+		if *called {
+			t.Fatal("next handler should NOT have been called")
+		}
+		assertJSONError(t, rec.Body.Bytes())
+	})
+
+	t.Run("valid token -> next called with user id in context", func(t *testing.T) {
+		store := newFakeSessionStore()
+		store.sessions["valid-token"] = &models.Session{
+			Token:  "valid-token",
+			UserID: "user-123",
+		}
+		rec, called, seenID := callMiddleware(t, store, "Bearer valid-token")
+		if rec.Code != http.StatusOK {
+			t.Fatalf("want 200, got %d", rec.Code)
+		}
+		if !*called {
+			t.Fatal("next handler SHOULD have been called")
+		}
+		if *seenID != "user-123" {
+			t.Fatalf("want user id 'user-123' in context, got %q", *seenID)
+		}
+	})
+
+	t.Run("store error -> 401 with JSON body", func(t *testing.T) {
+		store := newFakeSessionStore()
+		store.err = errors.New("boom")
+		rec, called, _ := callMiddleware(t, store, "Bearer valid-token")
+		if rec.Code != http.StatusUnauthorized {
+			t.Fatalf("want 401, got %d", rec.Code)
+		}
+		if *called {
+			t.Fatal("next handler should NOT have been called")
+		}
+		if ct := rec.Header().Get("Content-Type"); ct != "application/json" {
+			t.Fatalf("want Content-Type application/json, got %q", ct)
+		}
+		assertJSONError(t, rec.Body.Bytes())
+	})
+}