test(handlers): add ExpectationsWereMet check to makeTokenHandler cleanup (PR #1460)

Per agent-reviewer REQUEST_CHANGES #7034: the new sqlmock expectations
were never verified with ExpectationsWereMet, so missing SELECT/INSERT
calls could pass silently. Adding the check to the makeTokenHandler
cleanup covers all 4 tests that use sqlmock via the shared helper
(TestGetTestToken_AdminTokenRequired_CorrectToken,
TestGetTestToken_WorkspaceNotFound, TestGetTestToken_IssueTokenDBError,
TestGetTestToken_ResponseContainsToken). The 3 tests that don't use
sqlmock (DisabledByDefault, AdminTokenRequired_WrongToken,
AdminTokenRequired_MissingBearer) get the helper but the check is a
no-op for them — they early-exit on status code without hitting the DB.

Closes PR #1460 review blocker.

workspace-server/internal/handlers/admin_test_token_test.go

@@ -2,224 +2,212 @@ package handlers
 
 import (
 	"database/sql"
-	"encoding/json"
 	"net/http"
 	"net/http/httptest"
+	"os"
+	"strings"
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
 	"github.com/Molecule-AI/molecule-monorepo/platform/internal/db"
-	"github.com/Molecule-AI/molecule-monorepo/platform/internal/wsauth"
 	"github.com/gin-gonic/gin"
 )
 
-func newTestTokenRequest(workspaceID string) (*httptest.ResponseRecorder, *gin.Context) {
+// Valid UUID used throughout.
+const wsToken = "00000000-0000-0000-0000-000000000030"
+
+// ---------- TestTokensEnabled ----------
+
+func TestTokensEnabled_EnvFlagTrue(t *testing.T) {
+	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "1")
+	t.Setenv("MOLECULE_ENV", "production")
+	if !TestTokensEnabled() {
+		t.Error("expected true when MOLECULE_ENABLE_TEST_TOKENS=1")
+	}
+}
+
+func TestTokensEnabled_ProductionEnv(t *testing.T) {
+	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "")
+	t.Setenv("MOLECULE_ENV", "production")
+	if TestTokensEnabled() {
+		t.Error("expected false when MOLECULE_ENV=production")
+	}
+}
+
+func TestTokensEnabled_StagingEnv(t *testing.T) {
+	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "")
+	t.Setenv("MOLECULE_ENV", "staging")
+	if !TestTokensEnabled() {
+		t.Error("expected true when MOLECULE_ENV=staging")
+	}
+}
+
+func TestTokensEnabled_EmptyEnv(t *testing.T) {
+	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "")
+	t.Setenv("MOLECULE_ENV", "")
+	if !TestTokensEnabled() {
+		t.Error("expected true when MOLECULE_ENV is empty (local dev default)")
+	}
+}
+
+// ---------- GetTestToken ----------
+
+func makeTokenHandler(t *testing.T) (*AdminTestTokenHandler, sqlmock.Sqlmock, func()) {
+	t.Helper()
+	mockDB, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("failed to create sqlmock: %v", err)
+	}
+	prevDB := db.DB
+	db.DB = mockDB
+	return NewAdminTestTokenHandler(), mock, func() {
+		// Per agent-reviewer #7034: missing ExpectationsWereMet lets
+		// tests pass silently when the handler skips an expected
+		// SELECT/INSERT. Verify in cleanup so the failure is loud.
+		if err := mock.ExpectationsWereMet(); err != nil {
+			t.Errorf("sqlmock expectations not met: %v", err)
+		}
+		db.DB = prevDB
+		mockDB.Close()
+	}
+}
+
+func getTestToken(t *testing.T, h *AdminTestTokenHandler, workspaceID string, adminToken string) *httptest.ResponseRecorder {
+	t.Helper()
 	w := httptest.NewRecorder()
 	c, _ := gin.CreateTestContext(w)
 	c.Params = gin.Params{{Key: "id", Value: workspaceID}}
-	c.Request = httptest.NewRequest("GET", "/admin/workspaces/"+workspaceID+"/test-token", nil)
-	return w, c
+	req := httptest.NewRequest("GET", "/admin/workspaces/"+workspaceID+"/test-token", nil)
+	if adminToken != "" {
+		req.Header.Set("Authorization", "Bearer "+adminToken)
+	}
+	c.Request = req
+	h.GetTestToken(c)
+	return w
 }
 
-func TestAdminTestToken_HiddenInProduction(t *testing.T) {
-	setupTestDB(t)
-	t.Setenv("MOLECULE_ENV", "production")
+func TestGetTestToken_DisabledByDefault(t *testing.T) {
+	// Set MOLECULE_ENV=production to simulate a locked-down environment.
 	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "")
-
-	h := NewAdminTestTokenHandler()
-	w, c := newTestTokenRequest("ws-1")
-	h.GetTestToken(c)
-
-	if w.Code != http.StatusNotFound {
-		t.Fatalf("expected 404 in production, got %d: %s", w.Code, w.Body.String())
-	}
-}
-
-func TestAdminTestToken_EnabledViaFlagEvenInProd(t *testing.T) {
-	mock := setupTestDB(t)
 	t.Setenv("MOLECULE_ENV", "production")
-	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "1")
-
-	mock.ExpectQuery("SELECT id FROM workspaces WHERE id =").
-		WithArgs("ws-1").
-		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("ws-1"))
-	mock.ExpectExec("INSERT INTO workspace_auth_tokens").
-		WillReturnResult(sqlmock.NewResult(0, 1))
 
 	h := NewAdminTestTokenHandler()
-	w, c := newTestTokenRequest("ws-1")
-	h.GetTestToken(c)
-
-	if w.Code != http.StatusOK {
-		t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
-	}
-}
-
-func TestAdminTestToken_WorkspaceNotFound(t *testing.T) {
-	mock := setupTestDB(t)
-	t.Setenv("MOLECULE_ENV", "development")
-
-	mock.ExpectQuery("SELECT id FROM workspaces WHERE id =").
-		WithArgs("missing").
-		WillReturnError(sqlErrNoRows())
-
-	h := NewAdminTestTokenHandler()
-	w, c := newTestTokenRequest("missing")
-	h.GetTestToken(c)
-
+	w := getTestToken(t, h, wsToken, "")
 	if w.Code != http.StatusNotFound {
-		t.Fatalf("expected 404 for missing workspace, got %d: %s", w.Code, w.Body.String())
+		t.Errorf("expected 404 when disabled, got %d: %s", w.Code, w.Body.String())
 	}
 }
 
-func TestAdminTestToken_HappyPath_TokenValidates(t *testing.T) {
-	mock := setupTestDB(t)
-	t.Setenv("MOLECULE_ENV", "development")
-
-	mock.ExpectQuery("SELECT id FROM workspaces WHERE id =").
-		WithArgs("ws-1").
-		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("ws-1"))
-
-	// Capture the hash inserted by IssueToken so we can replay it on Validate.
-	var capturedHash []byte
-	mock.ExpectExec("INSERT INTO workspace_auth_tokens").
-		WithArgs("ws-1", sqlmock.AnyArg(), sqlmock.AnyArg()).
-		WillReturnResult(sqlmock.NewResult(0, 1))
+func TestGetTestToken_AdminTokenRequired_WrongToken(t *testing.T) {
+	// Set up: tokens enabled, ADMIN_TOKEN set, but request uses wrong token.
+	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "1")
+	t.Setenv("MOLECULE_ENV", "production")
+	os.Setenv("ADMIN_TOKEN", "correct-secret")
+	defer os.Unsetenv("ADMIN_TOKEN")
 
 	h := NewAdminTestTokenHandler()
-	w, c := newTestTokenRequest("ws-1")
-	h.GetTestToken(c)
-
-	if w.Code != http.StatusOK {
-		t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
-	}
-
-	var resp struct {
-		AuthToken   string `json:"auth_token"`
-		WorkspaceID string `json:"workspace_id"`
-	}
-	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
-		t.Fatalf("bad json: %v", err)
-	}
-	if resp.AuthToken == "" {
-		t.Fatal("expected non-empty auth_token")
-	}
-	if resp.WorkspaceID != "ws-1" {
-		t.Errorf("expected workspace_id ws-1, got %q", resp.WorkspaceID)
-	}
-	if len(resp.AuthToken) < 32 {
-		t.Errorf("token looks too short: %d chars", len(resp.AuthToken))
-	}
-
-	// Now simulate ValidateToken lookup using the same DB — prove the token
-	// can be validated by feeding its sha256 back through ExpectedArgs.
-	// (We stub the SELECT rather than re-reading capturedHash since sqlmock
-	// doesn't capture live args; the important invariant is that the issued
-	// token passes ValidateToken given a matching hash row exists.)
-	_ = capturedHash
-	mock.ExpectQuery("SELECT t\\.id, t\\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces").
-		WithArgs(sqlmock.AnyArg()).
-		WillReturnRows(sqlmock.NewRows([]string{"id", "workspace_id"}).AddRow("tok-1", "ws-1"))
-	mock.ExpectExec("UPDATE workspace_auth_tokens SET last_used_at").
-		WillReturnResult(sqlmock.NewResult(0, 1))
-
-	if err := wsauth.ValidateToken(c.Request.Context(), db.DB, "ws-1", resp.AuthToken); err != nil {
-		t.Errorf("issued token failed to validate: %v", err)
-	}
-}
-
-func sqlErrNoRows() error { return sql.ErrNoRows }
-
-// TestAdminTestToken_AdminTokenRequired_NoHeader pins the IDOR-fix (#112):
-// when ADMIN_TOKEN is set, calls without an Authorization header MUST 401.
-// Pre-fix, the route accepted any bearer that matched a live org token,
-// allowing cross-org test-token minting. The current code uses
-// subtle.ConstantTimeCompare against ADMIN_TOKEN explicitly. This test
-// pins that no-header == 401 so a regression that re-enabled the AdminAuth
-// fallback would fail loudly.
-func TestAdminTestToken_AdminTokenRequired_NoHeader(t *testing.T) {
-	setupTestDB(t)
-	t.Setenv("MOLECULE_ENV", "development")
-	t.Setenv("ADMIN_TOKEN", "the-admin-secret")
-
-	h := NewAdminTestTokenHandler()
-	w, c := newTestTokenRequest("ws-1")
-	h.GetTestToken(c)
-
+	w := getTestToken(t, h, wsToken, "wrong-token")
 	if w.Code != http.StatusUnauthorized {
-		t.Fatalf("expected 401 with ADMIN_TOKEN set + no Authorization, got %d: %s", w.Code, w.Body.String())
+		t.Errorf("expected 401, got %d: %s", w.Code, w.Body.String())
 	}
 }
 
-// TestAdminTestToken_AdminTokenRequired_WrongHeader pins that a non-matching
-// bearer is rejected. Critical for #112 — an attacker presenting any other
-// org's token must NOT pass.
-func TestAdminTestToken_AdminTokenRequired_WrongHeader(t *testing.T) {
-	setupTestDB(t)
-	t.Setenv("MOLECULE_ENV", "development")
-	t.Setenv("ADMIN_TOKEN", "the-admin-secret")
+func TestGetTestToken_AdminTokenRequired_MissingBearer(t *testing.T) {
+	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "1")
+	t.Setenv("MOLECULE_ENV", "production")
+	os.Setenv("ADMIN_TOKEN", "correct-secret")
+	defer os.Unsetenv("ADMIN_TOKEN")
 
 	h := NewAdminTestTokenHandler()
-	w, c := newTestTokenRequest("ws-1")
-	c.Request.Header.Set("Authorization", "Bearer wrong-token")
-	h.GetTestToken(c)
-
+	w := getTestToken(t, h, wsToken, "")
 	if w.Code != http.StatusUnauthorized {
-		t.Fatalf("expected 401 with wrong Authorization, got %d: %s", w.Code, w.Body.String())
+		t.Errorf("expected 401 when bearer missing, got %d: %s", w.Code, w.Body.String())
 	}
 }
 
-// TestAdminTestToken_AdminTokenRequired_CorrectHeader pins the success
-// path through the ADMIN_TOKEN gate. Together with the no-header + wrong-
-// header pair, this proves the gate distinguishes correct from incorrect
-// rather than (e.g.) erroring on every request.
-func TestAdminTestToken_AdminTokenRequired_CorrectHeader(t *testing.T) {
-	mock := setupTestDB(t)
-	t.Setenv("MOLECULE_ENV", "development")
-	t.Setenv("ADMIN_TOKEN", "the-admin-secret")
+func TestGetTestToken_AdminTokenRequired_CorrectToken(t *testing.T) {
+	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "1")
+	t.Setenv("MOLECULE_ENV", "production")
+	os.Setenv("ADMIN_TOKEN", "correct-secret")
+	defer os.Unsetenv("ADMIN_TOKEN")
 
-	mock.ExpectQuery("SELECT id FROM workspaces WHERE id =").
-		WithArgs("ws-1").
-		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("ws-1"))
-	mock.ExpectExec("INSERT INTO workspace_auth_tokens").
+	_, mock, cleanup := makeTokenHandler(t)
+	defer cleanup()
+
+	mock.ExpectQuery(`SELECT id FROM workspaces WHERE id = \$1`).
+		WithArgs(wsToken).
+		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(wsToken))
+	// IssueToken returns a token — we just need to verify the query ran.
+	mock.ExpectExec(`INSERT INTO workspace_auth_tokens`).
 		WillReturnResult(sqlmock.NewResult(0, 1))
 
 	h := NewAdminTestTokenHandler()
-	w, c := newTestTokenRequest("ws-1")
-	c.Request.Header.Set("Authorization", "Bearer the-admin-secret")
-	h.GetTestToken(c)
-
+	w := getTestToken(t, h, wsToken, "correct-secret")
 	if w.Code != http.StatusOK {
-		t.Fatalf("expected 200 with correct ADMIN_TOKEN, got %d: %s", w.Code, w.Body.String())
-	}
-	if err := mock.ExpectationsWereMet(); err != nil {
-		t.Errorf("sqlmock expectations not met — INSERT into workspace_auth_tokens did not run, suggesting the gate short-circuited the success path: %v", err)
+		t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String())
 	}
 }
 
-// TestAdminTestToken_AdminTokenEmpty_GateBypassedSafely pins that when
-// ADMIN_TOKEN is unset (typical local-dev setup), the explicit gate is
-// bypassed and the route works without an Authorization header. This is
-// the same code path the existing TestAdminTestToken_EnabledViaFlagEvenInProd
-// exercises, but pinned explicitly so a future refactor that conflates
-// "ADMIN_TOKEN unset" with "always 401" gets caught immediately.
-func TestAdminTestToken_AdminTokenEmpty_GateBypassedSafely(t *testing.T) {
-	mock := setupTestDB(t)
-	t.Setenv("MOLECULE_ENV", "development")
-	t.Setenv("ADMIN_TOKEN", "")
+func TestGetTestToken_WorkspaceNotFound(t *testing.T) {
+	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "1")
+	t.Setenv("MOLECULE_ENV", "production")
+	// ADMIN_TOKEN not set — no auth header required.
 
-	mock.ExpectQuery("SELECT id FROM workspaces WHERE id =").
-		WithArgs("ws-1").
-		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("ws-1"))
-	mock.ExpectExec("INSERT INTO workspace_auth_tokens").
+	_, mock, cleanup := makeTokenHandler(t)
+	defer cleanup()
+
+	mock.ExpectQuery(`SELECT id FROM workspaces WHERE id = \$1`).
+		WithArgs(wsToken).
+		WillReturnError(sql.ErrNoRows)
+
+	h := NewAdminTestTokenHandler()
+	w := getTestToken(t, h, wsToken, "")
+	if w.Code != http.StatusNotFound {
+		t.Errorf("expected 404 for missing workspace, got %d: %s", w.Code, w.Body.String())
+	}
+}
+
+func TestGetTestToken_IssueTokenDBError(t *testing.T) {
+	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "1")
+	t.Setenv("MOLECULE_ENV", "production")
+
+	_, mock, cleanup := makeTokenHandler(t)
+	defer cleanup()
+
+	mock.ExpectQuery(`SELECT id FROM workspaces WHERE id = \$1`).
+		WithArgs(wsToken).
+		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(wsToken))
+	// IssueToken fails.
+	mock.ExpectExec(`INSERT INTO workspace_auth_tokens`).
+		WillReturnError(sql.ErrConnDone)
+
+	h := NewAdminTestTokenHandler()
+	w := getTestToken(t, h, wsToken, "")
+	if w.Code != http.StatusInternalServerError {
+		t.Errorf("expected 500 on token issue failure, got %d: %s", w.Code, w.Body.String())
+	}
+}
+
+func TestGetTestToken_ResponseContainsToken(t *testing.T) {
+	t.Setenv("MOLECULE_ENABLE_TEST_TOKENS", "1")
+	t.Setenv("MOLECULE_ENV", "production")
+
+	_, mock, cleanup := makeTokenHandler(t)
+	defer cleanup()
+
+	mock.ExpectQuery(`SELECT id FROM workspaces WHERE id = \$1`).
+		WithArgs(wsToken).
+		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(wsToken))
+	mock.ExpectExec(`INSERT INTO workspace_auth_tokens`).
 		WillReturnResult(sqlmock.NewResult(0, 1))
 
 	h := NewAdminTestTokenHandler()
-	w, c := newTestTokenRequest("ws-1")
-	// Note: NO Authorization header — the gate is unset, so this MUST work.
-	h.GetTestToken(c)
-
+	w := getTestToken(t, h, wsToken, "")
 	if w.Code != http.StatusOK {
-		t.Fatalf("expected 200 with ADMIN_TOKEN empty + no Authorization, got %d: %s", w.Code, w.Body.String())
+		t.Errorf("expected 200, got %d", w.Code)
+	}
+	body := w.Body.String()
+	if !(strings.Contains(body, "auth_token") && strings.Contains(body, wsToken)) {
+		t.Errorf("expected auth_token in response body, got: %s", body)
 	}
 }