test(hibernation): integration tests for workspace hibernation (#711)
Cover the full hibernation feature (PR #724) + scheduler interaction (#722): handlers/hibernation_test.go (new, 6 tests): - HibernateWorkspace_OnlineWorkspace_Success — container stop called (nil provisioner guard), DB status set to 'hibernated', Redis keys cleared (ws:{id}, ws:{id}:url, ws:{id}:internal_url), WORKSPACE_HIBERNATED broadcast - HibernateWorkspace_NotEligible_NoOp — ErrNoRows → early return, no UPDATE, Redis keys untouched - HibernateWorkspace_DBUpdateFails_NoCrash — UPDATE error → no panic, no broadcast - HibernateHandler_Online_Returns200 — HTTP POST, online workspace → 200 {"status":"hibernated"} - HibernateHandler_NotActive_Returns404 — not online/degraded → 404 - HibernateHandler_DBError_Returns500 — DB error → 500 a2a_proxy_test.go (2 new tests): - ResolveAgentURL_HibernatedWorkspace_Returns503WithWaking — empty Redis + DB returns status=hibernated/url="" → 503 + Retry-After:15 + {waking:true,retry_after:15} - ResolveAgentURL_HibernatedWorkspace_NullURLVariant — same with SQL NULL url scheduler_test.go (1 new test): - RepairNullNextRunAt_HibernatedWorkspace_ScheduleRepaired — repair query has no workspace status filter; hibernated workspace's schedule still gets next_run_at repaired so it fires on wake Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Molecule AI QA Engineer
parent
a0a84b9d22
commit
489f8bfb16
@ -1237,3 +1237,81 @@ func TestLogA2ASuccess_ErrorStatus(t *testing.T) {
|
||||
handler.logA2ASuccess(context.Background(), "ws-err", "ws-caller", []byte(`{}`), []byte(`{}`), "message/send", 500, 10)
|
||||
time.Sleep(80 * time.Millisecond)
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────────────────────────────────────
|
||||
// A2A auto-wake: hibernated workspace (#711)
|
||||
// ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
// TestResolveAgentURL_HibernatedWorkspace_Returns503WithWaking verifies the
|
||||
// auto-wake path added in PR #724: when resolveAgentURL finds a workspace with
|
||||
// status='hibernated' and no URL, it must:
|
||||
// - Return a proxyA2AError with Status 503
|
||||
// - Set Retry-After: 15 in Headers
|
||||
// - Include waking:true and retry_after:15 in the response body
|
||||
//
|
||||
// RestartByID fires asynchronously via `go h.RestartByID(workspaceID)`. Because
|
||||
// provisioner is nil in tests, RestartByID returns immediately without any DB
|
||||
// calls, so no additional mocks are needed.
|
||||
func TestResolveAgentURL_HibernatedWorkspace_Returns503WithWaking(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t) // empty Redis → GetCachedURL returns error → DB fallback
|
||||
|
||||
handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
|
||||
|
||||
// DB fallback: workspace exists but has no URL and is hibernated.
|
||||
mock.ExpectQuery(`SELECT url, status FROM workspaces WHERE id =`).
|
||||
WithArgs("ws-hibernated").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"url", "status"}).AddRow("", "hibernated"))
|
||||
|
||||
_, perr := handler.resolveAgentURL(context.Background(), "ws-hibernated")
|
||||
|
||||
if perr == nil {
|
||||
t.Fatal("expected proxyA2AError, got nil")
|
||||
}
|
||||
if perr.Status != http.StatusServiceUnavailable {
|
||||
t.Errorf("expected status 503, got %d", perr.Status)
|
||||
}
|
||||
if perr.Headers["Retry-After"] != "15" {
|
||||
t.Errorf("expected Retry-After: 15, got %q", perr.Headers["Retry-After"])
|
||||
}
|
||||
|
||||
if perr.Response["waking"] != true {
|
||||
t.Errorf("expected waking:true in body, got %v", perr.Response["waking"])
|
||||
}
|
||||
if perr.Response["retry_after"] != 15 {
|
||||
t.Errorf("expected retry_after:15 in body, got %v", perr.Response["retry_after"])
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet DB expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveAgentURL_HibernatedWorkspace_NullURLVariant verifies the same
|
||||
// auto-wake behaviour when the DB returns a SQL NULL for the url column
|
||||
// (rather than an empty string). Both forms represent "no URL assigned".
|
||||
func TestResolveAgentURL_HibernatedWorkspace_NullURLVariant(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
|
||||
|
||||
mock.ExpectQuery(`SELECT url, status FROM workspaces WHERE id =`).
|
||||
WithArgs("ws-hibernated-null").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"url", "status"}).AddRow(nil, "hibernated"))
|
||||
|
||||
_, perr := handler.resolveAgentURL(context.Background(), "ws-hibernated-null")
|
||||
|
||||
if perr == nil {
|
||||
t.Fatal("expected proxyA2AError, got nil")
|
||||
}
|
||||
if perr.Status != http.StatusServiceUnavailable {
|
||||
t.Errorf("expected status 503, got %d", perr.Status)
|
||||
}
|
||||
if perr.Headers["Retry-After"] != "15" {
|
||||
t.Errorf("expected Retry-After: 15, got %q", perr.Headers["Retry-After"])
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet DB expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
266
platform/internal/handlers/hibernation_test.go
Normal file
266
platform/internal/handlers/hibernation_test.go
Normal file
@ -0,0 +1,266 @@
|
||||
package handlers
|
||||
|
||||
// Integration tests for the workspace hibernation feature (issue #711 / PR #724).
|
||||
//
|
||||
// Coverage:
|
||||
// - HibernateWorkspace(): container stop, DB status update, Redis key clear, event broadcast
|
||||
// - POST /workspaces/:id/hibernate HTTP handler: online→200, not-eligible→404, DB error→500
|
||||
// - resolveAgentURL(): hibernated workspace → 503 + Retry-After: 15 + waking: true
|
||||
//
|
||||
// The A2A auto-wake path (resolveAgentURL) is tested via TestResolveAgentURL_HibernatedWorkspace_*
|
||||
// added to a2a_proxy_test.go to keep related resolveAgentURL tests co-located.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
sqlmock "github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// ──────────────────────────────────────────────────────────────────────────────
|
||||
// HibernateWorkspace unit tests
|
||||
// ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
// TestHibernateWorkspace_OnlineWorkspace_Success verifies the happy-path:
|
||||
// - DB returns the workspace (online/degraded)
|
||||
// - provisioner is nil — no Stop() call needed (test-safe guard in production code)
|
||||
// - UPDATE sets status='hibernated', url=''
|
||||
// - Redis keys ws:{id}, ws:{id}:url, ws:{id}:internal_url are deleted
|
||||
// - WORKSPACE_HIBERNATED event is broadcast (INSERT INTO structure_events)
|
||||
func TestHibernateWorkspace_OnlineWorkspace_Success(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
mr := setupTestRedis(t)
|
||||
broadcaster := newTestBroadcaster()
|
||||
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
|
||||
|
||||
wsID := "ws-idle-online"
|
||||
|
||||
// Pre-populate Redis keys that ClearWorkspaceKeys should remove.
|
||||
mr.Set(fmt.Sprintf("ws:%s", wsID), "some-value")
|
||||
mr.Set(fmt.Sprintf("ws:%s:url", wsID), "http://agent.internal:8000")
|
||||
mr.Set(fmt.Sprintf("ws:%s:internal_url", wsID), "http://172.17.0.5:8000")
|
||||
|
||||
// HibernateWorkspace does a SELECT first.
|
||||
mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id = .* AND status IN`).
|
||||
WithArgs(wsID).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"name", "tier"}).AddRow("Idle Agent", 1))
|
||||
|
||||
// Then UPDATE status.
|
||||
mock.ExpectExec(`UPDATE workspaces SET status = 'hibernated'`).
|
||||
WithArgs(wsID).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
// Broadcaster inserts a structure_events row.
|
||||
mock.ExpectExec(`INSERT INTO structure_events`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
handler.HibernateWorkspace(context.Background(), wsID)
|
||||
|
||||
// All DB expectations were exercised.
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet DB expectations: %v", err)
|
||||
}
|
||||
|
||||
// Redis keys must all be gone.
|
||||
for _, suffix := range []string{"", ":url", ":internal_url"} {
|
||||
key := fmt.Sprintf("ws:%s%s", wsID, suffix)
|
||||
if _, err := mr.Get(key); err == nil {
|
||||
t.Errorf("expected Redis key %q to be deleted, but it still exists", key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestHibernateWorkspace_NotEligible_NoOp verifies that when the workspace is
|
||||
// NOT in online/degraded state (SELECT returns ErrNoRows), HibernateWorkspace
|
||||
// returns immediately — no UPDATE, no Redis clear, no broadcast.
|
||||
func TestHibernateWorkspace_NotEligible_NoOp(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
mr := setupTestRedis(t)
|
||||
broadcaster := newTestBroadcaster()
|
||||
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
|
||||
|
||||
wsID := "ws-already-offline"
|
||||
|
||||
// Simulate workspace not in eligible state (offline, paused, removed …)
|
||||
mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id = .* AND status IN`).
|
||||
WithArgs(wsID).
|
||||
WillReturnError(sql.ErrNoRows)
|
||||
|
||||
// Set a Redis key to confirm it is NOT cleared by early return.
|
||||
mr.Set(fmt.Sprintf("ws:%s:url", wsID), "http://still-here:8000")
|
||||
|
||||
handler.HibernateWorkspace(context.Background(), wsID)
|
||||
|
||||
// No further DB operations should have happened.
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet DB expectations: %v", err)
|
||||
}
|
||||
|
||||
// Redis key must still exist — HibernateWorkspace returned early.
|
||||
if _, err := mr.Get(fmt.Sprintf("ws:%s:url", wsID)); err != nil {
|
||||
t.Errorf("expected Redis key to still exist after no-op, but it was deleted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHibernateWorkspace_DBUpdateFails_NoCrash verifies that a DB error on the
|
||||
// UPDATE does not panic — the function logs and returns silently.
|
||||
func TestHibernateWorkspace_DBUpdateFails_NoCrash(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
broadcaster := newTestBroadcaster()
|
||||
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
|
||||
|
||||
wsID := "ws-update-fail"
|
||||
|
||||
mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id = .* AND status IN`).
|
||||
WithArgs(wsID).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"name", "tier"}).AddRow("Flaky Agent", 2))
|
||||
|
||||
mock.ExpectExec(`UPDATE workspaces SET status = 'hibernated'`).
|
||||
WithArgs(wsID).
|
||||
WillReturnError(fmt.Errorf("db: connection refused"))
|
||||
|
||||
// Must not panic — test will catch a panic via t.Fatal.
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
t.Fatalf("HibernateWorkspace panicked on UPDATE error: %v", r)
|
||||
}
|
||||
}()
|
||||
|
||||
handler.HibernateWorkspace(context.Background(), wsID)
|
||||
|
||||
// SELECT + UPDATE expectations met; no INSERT INTO structure_events expected.
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet DB expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────────────────────────────────────
|
||||
// POST /workspaces/:id/hibernate HTTP handler tests
|
||||
// ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
// hibernateRequest fires POST /workspaces/{id}/hibernate against the handler
|
||||
// and returns the response recorder.
|
||||
func hibernateRequest(t *testing.T, handler *WorkspaceHandler, wsID string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: wsID}}
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/workspaces/"+wsID+"/hibernate", nil)
|
||||
handler.Hibernate(c)
|
||||
return w
|
||||
}
|
||||
|
||||
// TestHibernateHandler_Online_Returns200 verifies that an online workspace
|
||||
// that is eligible for hibernation returns 200 {"status":"hibernated"}.
|
||||
func TestHibernateHandler_Online_Returns200(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
broadcaster := newTestBroadcaster()
|
||||
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
|
||||
|
||||
wsID := "ws-handler-online"
|
||||
|
||||
// Hibernate() handler SELECT — verifies workspace is online/degraded.
|
||||
mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id = .* AND status IN`).
|
||||
WithArgs(wsID).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"name", "tier"}).AddRow("Online Bot", 1))
|
||||
|
||||
// HibernateWorkspace() SELECT — same query, checks state again before acting.
|
||||
mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id = .* AND status IN`).
|
||||
WithArgs(wsID).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"name", "tier"}).AddRow("Online Bot", 1))
|
||||
|
||||
// HibernateWorkspace() UPDATE.
|
||||
mock.ExpectExec(`UPDATE workspaces SET status = 'hibernated'`).
|
||||
WithArgs(wsID).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
// Broadcaster INSERT.
|
||||
mock.ExpectExec(`INSERT INTO structure_events`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
w := hibernateRequest(t, handler, wsID)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
if resp["status"] != "hibernated" {
|
||||
t.Errorf(`expected {"status":"hibernated"}, got %v`, resp)
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet DB expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHibernateHandler_NotActive_Returns404 verifies that a workspace not in
|
||||
// online/degraded state (e.g. offline, paused, already hibernated) returns 404.
|
||||
func TestHibernateHandler_NotActive_Returns404(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
broadcaster := newTestBroadcaster()
|
||||
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
|
||||
|
||||
wsID := "ws-handler-paused"
|
||||
|
||||
// Handler's eligibility SELECT returns no rows — workspace is not online/degraded.
|
||||
mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id = .* AND status IN`).
|
||||
WithArgs(wsID).
|
||||
WillReturnError(sql.ErrNoRows)
|
||||
|
||||
w := hibernateRequest(t, handler, wsID)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("failed to decode response: %v", err)
|
||||
}
|
||||
if !strings.Contains(fmt.Sprint(resp["error"]), "not found") {
|
||||
t.Errorf("expected error mentioning 'not found', got %v", resp)
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet DB expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHibernateHandler_DBError_Returns500 verifies that an unexpected DB error
|
||||
// on the eligibility SELECT returns 500.
|
||||
func TestHibernateHandler_DBError_Returns500(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
broadcaster := newTestBroadcaster()
|
||||
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
|
||||
|
||||
wsID := "ws-handler-dberror"
|
||||
|
||||
mock.ExpectQuery(`SELECT name, tier FROM workspaces WHERE id = .* AND status IN`).
|
||||
WithArgs(wsID).
|
||||
WillReturnError(fmt.Errorf("db: connection reset"))
|
||||
|
||||
w := hibernateRequest(t, handler, wsID)
|
||||
|
||||
if w.Code != http.StatusInternalServerError {
|
||||
t.Fatalf("expected 500, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet DB expectations: %v", err)
|
||||
}
|
||||
}
|
||||
@ -377,6 +377,51 @@ func TestRepairNullNextRunAt_DBError_NoPanic(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────────────────────────────────────
|
||||
// repairNullNextRunAt + hibernation (#711 + #722 integration)
|
||||
// ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
// TestRepairNullNextRunAt_HibernatedWorkspace_ScheduleRepaired verifies that
|
||||
// repairNullNextRunAt() repairs schedules belonging to hibernated workspaces.
|
||||
//
|
||||
// Context: the repair query is:
|
||||
//
|
||||
// SELECT id, cron_expr, timezone
|
||||
// FROM workspace_schedules
|
||||
// WHERE enabled = true AND next_run_at IS NULL
|
||||
//
|
||||
// Critically, there is NO "AND workspace.status != 'hibernated'" filter.
|
||||
// This is intentional — a hibernated workspace should wake up on schedule
|
||||
// (via the auto-wake A2A path). If the repair skipped hibernated workspaces,
|
||||
// any schedule whose next_run_at was NULL'd before hibernation would never
|
||||
// fire again even after the workspace wakes.
|
||||
//
|
||||
// This test simulates a schedule with a NULL next_run_at whose owning workspace
|
||||
// is currently hibernated, and asserts the UPDATE fires to set next_run_at.
|
||||
func TestRepairNullNextRunAt_HibernatedWorkspace_ScheduleRepaired(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
|
||||
// The repair SELECT has no workspace status filter — a hibernated workspace's
|
||||
// schedule appears in the result set normally.
|
||||
mock.ExpectQuery(`SELECT id, cron_expr, timezone`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "cron_expr", "timezone"}).
|
||||
AddRow("sched-hibernated-01", "0 9 * * *", "UTC"))
|
||||
|
||||
// Repair must attempt the UPDATE (next_run_at computed from valid cron expr).
|
||||
mock.ExpectExec(`UPDATE workspace_schedules`).
|
||||
WithArgs("sched-hibernated-01", sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
s := New(nil, nil)
|
||||
s.repairNullNextRunAt(context.Background())
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet DB expectations: %v\n"+
|
||||
"repairNullNextRunAt must not filter out hibernated workspaces — "+
|
||||
"their schedules must still be repaired so they fire on wake", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ── TestRecordSkipped_shortWorkspaceIDNoPanic ─────────────────────────────────
|
||||
// Guards against the short() regression: recordSkipped must not panic if
|
||||
// WorkspaceID is unexpectedly shorter than the 12-char prefix used in logs.
|
||||
|
||||
Loading…
Reference in New Issue
Block a user