package server import ( "context" "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/mem9-ai/drive9/pkg/meta" "github.com/mem9-ai/drive9/pkg/metrics" ) func TestRequestRoute(t *testing.T) { t.Parallel() tests := []struct { method string target string want string }{ {method: http.MethodGet, target: "/healthz", want: "/healthz"}, {method: http.MethodGet, target: "/metrics", want: "/metrics"}, {method: http.MethodPost, target: "/v1/provision", want: "/v1/provision"}, {method: http.MethodGet, target: "/v1/status", want: "/v1/status"}, {method: http.MethodGet, target: "/v1/tokens", want: "/v1/tokens/*?list"}, {method: http.MethodPost, target: "/v1/tokens", want: "/v1/tokens/*?issue"}, {method: http.MethodGet, target: "/v1/tokens/key1", want: "/v1/tokens/*?unmatched"}, {method: http.MethodDelete, target: "/v1/tokens/key1", want: "/v1/tokens/*?revoke"}, {method: http.MethodDelete, target: "/v1/tokens/revoke", want: "/v1/tokens/*?unmatched"}, {method: http.MethodDelete, target: "/v1/tokens/generate", want: "/v1/tokens/*?unmatched"}, {method: http.MethodPost, target: "/v1/tokens/generate", want: "/v1/tokens/*?generate"}, {method: http.MethodPost, target: "/v1/tokens/refresh", want: "/v1/tokens/*?refresh"}, {method: http.MethodPost, target: "/v1/tokens/revoke", want: "/v1/tokens/*?revoke_by_key"}, {method: http.MethodPost, target: "/v1/tokens/key1/activate", want: "/v1/tokens/*?activate"}, {method: http.MethodPost, target: "/v1/tokens/key1/deactivate", want: "/v1/tokens/*?deactivate"}, {method: http.MethodPost, target: "/v1/sql", want: "/v1/sql"}, {method: http.MethodGet, target: "/v1/events", want: "/v1/events"}, {method: http.MethodPost, target: "/v1/fork", want: "/v1/fork"}, {method: http.MethodGet, target: "/v1/journals", want: "/v1/journals"}, {method: http.MethodGet, target: "/v1/journals/j1", want: "/v1/journals/*"}, {method: http.MethodPost, target: "/v1/journals/j1/entries", want: "/v1/journals/*?entries"}, {method: http.MethodGet, target: "/v1/journals/j1/entries", want: "/v1/journals/*?entries"}, {method: http.MethodGet, target: "/v1/journals/j1/verify", want: "/v1/journals/*?verify"}, {method: http.MethodGet, target: "/v1/journals/j1/bogus", want: "/v1/journals/*"}, {method: http.MethodGet, target: "/v1/journal-entries", want: "/v1/journal-entries"}, {method: http.MethodGet, target: "/v1/git-workspaces", want: "/v1/git-workspaces"}, {method: http.MethodGet, target: "/v1/git-workspaces/w1", want: "/v1/git-workspaces/*"}, {method: http.MethodGet, target: "/v1/git-workspaces/w1/tree", want: "/v1/git-workspaces/*?tree"}, {method: http.MethodPost, target: "/v1/git-workspaces/w1/git-state", want: "/v1/git-workspaces/*?git-state"}, {method: http.MethodGet, target: "/v1/git-workspaces/w1/object-packs", want: "/v1/git-workspaces/*?object-packs"}, {method: http.MethodGet, target: "/v1/git-workspaces/w1/object-packs/p1", want: "/v1/git-workspaces/*?object-packs"}, {method: http.MethodGet, target: "/v1/git-workspaces/w1/overlay", want: "/v1/git-workspaces/*?overlay"}, {method: http.MethodGet, target: "/v1/git-workspaces/w1/bogus", want: "/v1/git-workspaces/*"}, {method: http.MethodGet, target: "/v1/layers", want: "/v1/layers"}, {method: http.MethodGet, target: "/v1/layers/l1", want: "/v1/layers/*?get"}, {method: http.MethodGet, target: "/v1/layers/l1/diff", want: "/v1/layers/*?diff"}, {method: http.MethodPost, target: "/v1/layers/l1/fork", want: "/v1/layers/*?fork"}, {method: http.MethodGet, target: "/v1/layers/l1/chain", want: "/v1/layers/*?chain"}, {method: http.MethodPost, target: "/v1/layers/l1/commit", want: "/v1/layers/*?commit"}, {method: http.MethodGet, target: "/v1/layers/l1/entries", want: "/v1/layers/*?entries_get"}, {method: http.MethodPost, target: "/v1/layers/l1/entries", want: "/v1/layers/*?entries_post"}, {method: http.MethodPut, target: "/v1/layers/l1/entries", want: "/v1/layers/*?entries_put"}, {method: http.MethodGet, target: "/v1/layers/l1/objects", want: "/v1/layers/*?objects_get"}, {method: http.MethodPost, target: "/v1/layers/l1/objects", want: "/v1/layers/*?objects_post"}, {method: http.MethodPut, target: "/v1/layers/l1/objects", want: "/v1/layers/*?objects_put"}, {method: http.MethodGet, target: "/v1/layers/l1/events", want: "/v1/layers/*?events"}, {method: http.MethodPost, target: "/v1/layers/l1/rollback", want: "/v1/layers/*?rollback"}, {method: http.MethodPost, target: "/v1/layers/l1/checkpoints", want: "/v1/layers/*?checkpoint"}, {method: http.MethodDelete, target: "/v1/layers/l1", want: "/v1/layers/*?delete"}, {method: http.MethodPost, target: "/v1/layers/l1/uploads/initiate", want: "/v1/layers/*?upload_initiate"}, {method: http.MethodPost, target: "/v1/layers/l1/uploads/u1/presign", want: "/v1/layers/*?upload_presign"}, {method: http.MethodPost, target: "/v1/layers/l1/uploads/u1/presign-batch", want: "/v1/layers/*?upload_presign_batch"}, {method: http.MethodPost, target: "/v1/layers/l1/uploads/u1/complete", want: "/v1/layers/*?upload_complete"}, {method: http.MethodPost, target: "/v1/layers/l1/uploads/u1/abort", want: "/v1/layers/*?upload_abort"}, {method: http.MethodGet, target: "/v1/layers/l1/bogus", want: "/v1/layers/*?other"}, {method: http.MethodDelete, target: "/v1/layers/l1/bogus", want: "/v1/layers/*?other"}, // Slash variants 404 in handleFSLayerObject and must stay out of the // real operation buckets. {method: http.MethodGet, target: "/v1/layers/l1/diff/", want: "/v1/layers/*?other"}, {method: http.MethodGet, target: "/v1/layers/l1//diff", want: "/v1/layers/*?other"}, {method: http.MethodPost, target: "/v1/layers/l1/commit/", want: "/v1/layers/*?other"}, {method: http.MethodPost, target: "/v1/layers/l1/uploads/initiate/", want: "/v1/layers/*?other"}, {method: http.MethodGet, target: "/v1/layer-checkpoints/c1", want: "/v1/layer-checkpoints/*"}, {method: http.MethodGet, target: "/v1/admin/tenant-pool", want: "/v1/admin/tenant-pool"}, {method: http.MethodGet, target: "/v1/admin/tenants", want: "/v1/admin/tenants"}, {method: http.MethodGet, target: "/v1/admin/tenants/t1", want: "/v1/admin/tenants/*"}, {method: http.MethodPost, target: "/v1/admin/tenants/t1/quota", want: "/v1/admin/tenants/*?quota"}, {method: http.MethodPost, target: "/v1/admin/tenants/t1/quota/extra", want: "/v1/admin/tenants/*"}, {method: http.MethodPost, target: "/v1/admin/tenants//quota", want: "/v1/admin/tenants/*"}, {method: http.MethodGet, target: "/v1/admin/tenants/t1/extract-config/image", want: "/v1/admin/tenants/*?extract-config"}, {method: http.MethodPut, target: "/v1/admin/tenants/t1/extract-config/image", want: "/v1/admin/tenants/*?extract-config"}, {method: http.MethodGet, target: "/v1/auth/slock/login", want: "/v1/auth/slock/login"}, {method: http.MethodGet, target: "/v1/auth/slock/callback", want: "/v1/auth/slock/callback"}, {method: http.MethodPost, target: "/v1/migration/events", want: "/v1/migration/events"}, // FS wildcard route is split by the bounded action enum so // read/list/grep/find and POST actions stay in separate buckets. {method: http.MethodGet, target: "/v1/fs/doc.txt", want: "/v1/fs/*?read"}, {method: http.MethodGet, target: "/v1/fs/dir/?list=1", want: "/v1/fs/*?list"}, {method: http.MethodGet, target: "/v1/fs/dir/?grep=foo&limit=20", want: "/v1/fs/*?grep"}, {method: http.MethodGet, target: "/v1/fs/dir/?find=&name=*.txt", want: "/v1/fs/*?find"}, {method: http.MethodGet, target: "/v1/fs/doc.txt?stat", want: "/v1/fs/*?stat"}, {method: http.MethodHead, target: "/v1/fs/doc.txt", want: "/v1/fs/*?stat"}, {method: http.MethodPut, target: "/v1/fs/doc.txt", want: "/v1/fs/*?write"}, {method: http.MethodPatch, target: "/v1/fs/doc.txt", want: "/v1/fs/*?patch"}, {method: http.MethodDelete, target: "/v1/fs/doc.txt", want: "/v1/fs/*?delete"}, {method: http.MethodPost, target: "/v1/fs/dir/?mkdir", want: "/v1/fs/*?mkdir"}, {method: http.MethodPost, target: "/v1/fs/b.txt?copy", want: "/v1/fs/*?copy"}, {method: http.MethodPost, target: "/v1/fs/a.txt?append-log", want: "/v1/fs/*?append_log"}, {method: http.MethodPost, target: "/v1/fs/a.txt?append", want: "/v1/fs/*?append"}, {method: http.MethodPost, target: "/v1/fs/a.txt?rename", want: "/v1/fs/*?rename"}, {method: http.MethodPost, target: "/v1/fs/a.txt?chmod", want: "/v1/fs/*?chmod"}, {method: http.MethodPost, target: "/v1/fs/a.txt?create", want: "/v1/fs/*?create"}, {method: http.MethodPost, target: "/v1/fs/a.txt?symlink", want: "/v1/fs/*?symlink"}, {method: http.MethodPost, target: "/v1/fs/a.txt?hardlink=1", want: "/v1/fs/*?hardlink"}, {method: http.MethodPost, target: "/v1/fs/doc.txt", want: "/v1/fs/*?post"}, {method: http.MethodPost, target: "/v1/fs:batch-stat", want: "/v1/fs:batch-stat"}, {method: http.MethodPost, target: "/v1/fs:batch-read-small", want: "/v1/fs:batch-read-small"}, {method: http.MethodPost, target: "/v1/fs:batch-write", want: "/v1/fs:batch-write"}, {method: http.MethodGet, target: "/v1/uploads", want: "/v1/uploads"}, // POST /v1/uploads/initiate is a live alias for POST /v1/uploads. {method: http.MethodPost, target: "/v1/uploads/initiate", want: "/v1/uploads/*?initiate"}, {method: http.MethodGet, target: "/v1/uploads/initiate", want: "/v1/uploads/*?other"}, // Upload wildcards are split by their bounded action classifiers. {method: http.MethodPost, target: "/v1/uploads/u1/complete", want: "/v1/uploads/*?complete"}, {method: http.MethodPost, target: "/v1/uploads/u1/resume", want: "/v1/uploads/*?resume"}, {method: http.MethodDelete, target: "/v1/uploads/u1", want: "/v1/uploads/*?abort"}, {method: http.MethodGet, target: "/v1/uploads/u1/unknown", want: "/v1/uploads/*?other"}, {method: http.MethodPost, target: "/v2/uploads/initiate", want: "/v2/uploads/*?initiate"}, {method: http.MethodPost, target: "/v2/uploads/u1/presign", want: "/v2/uploads/*?presign"}, {method: http.MethodPost, target: "/v2/uploads/u1/presign-batch", want: "/v2/uploads/*?presign_batch"}, {method: http.MethodPost, target: "/v2/uploads/u1/complete", want: "/v2/uploads/*?complete"}, {method: http.MethodPost, target: "/v2/uploads/u1/abort", want: "/v2/uploads/*?abort"}, {method: http.MethodPost, target: "/v2/uploads/u1/parts", want: "/v2/uploads/*?other"}, {method: http.MethodGet, target: "/v1/vault/secrets", want: "/v1/vault/secrets/*"}, {method: http.MethodGet, target: "/v1/vault/secrets/db-prod", want: "/v1/vault/secrets/*"}, {method: http.MethodGet, target: "/v1/vault/tokens", want: "/v1/vault/tokens/*"}, {method: http.MethodGet, target: "/v1/vault/grants/g1", want: "/v1/vault/grants/*"}, {method: http.MethodGet, target: "/v1/vault/audit", want: "/v1/vault/audit"}, {method: http.MethodGet, target: "/v1/vault/read", want: "/v1/vault/read/*"}, {method: http.MethodGet, target: "/v1/vault/read/secret/path", want: "/v1/vault/read/*"}, {method: http.MethodPut, target: "/s3/tenant-a/object", want: "/s3/*?put"}, {method: http.MethodPut, target: "/s3/tenant-a/upload/u1", want: "/s3/*?upload_part"}, {method: http.MethodGet, target: "/s3/tenant-a/objects/o1", want: "/s3/*?get_object"}, {method: http.MethodGet, target: "/unknown", want: "other"}, } for _, tc := range tests { r := httptest.NewRequest(tc.method, tc.target, nil) if got := requestRoute(r); got != tc.want { t.Fatalf("requestRoute(%s %q) = %q, want %q", tc.method, tc.target, got, tc.want) } } } func TestClassifyTenantRequest(t *testing.T) { t.Parallel() tests := []struct { method string path string want tenantRequestClass }{ {method: http.MethodGet, path: "/v1/fs/doc.txt", want: tenantRequestClass{surface: "fs", action: "read"}}, {method: http.MethodGet, path: "/v1/fs/dir/?list=1", want: tenantRequestClass{surface: "fs", action: "list"}}, {method: http.MethodGet, path: "/v1/fs:batch-stat", want: tenantRequestClass{surface: "fs", action: "batch_stat"}}, {method: http.MethodPost, path: "/v1/fs:batch-read-small", want: tenantRequestClass{surface: "fs", action: "batch_read_small"}}, {method: http.MethodPost, path: "/v1/fs:batch-write", want: tenantRequestClass{surface: "fs", action: "batch_write"}}, {method: http.MethodPost, path: "/v1/fs/large.bin?append=1", want: tenantRequestClass{surface: "fs", action: "append"}}, {method: http.MethodPost, path: "/v1/fs/large.bin?append-log=1", want: tenantRequestClass{surface: "fs", action: "append_log"}}, {method: http.MethodPost, path: "/v1/uploads", want: tenantRequestClass{surface: "upload", action: "initiate"}}, {method: http.MethodGet, path: "/v1/uploads", want: tenantRequestClass{surface: "upload", action: "list"}}, {method: http.MethodPost, path: "/v1/uploads/u1/complete", want: tenantRequestClass{surface: "upload", action: "complete"}}, {method: http.MethodDelete, path: "/v1/uploads/u1", want: tenantRequestClass{surface: "upload", action: "abort"}}, {method: http.MethodPost, path: "/v1/uploads/u1/random-user-input", want: tenantRequestClass{surface: "upload", action: "other"}}, {method: http.MethodPost, path: "/v2/uploads/initiate", want: tenantRequestClass{surface: "upload", action: "initiate"}}, {method: http.MethodPost, path: "/v2/uploads/u1/presign-batch", want: tenantRequestClass{surface: "upload", action: "presign_batch"}}, {method: http.MethodPost, path: "/v2/uploads/u1/random-user-input", want: tenantRequestClass{surface: "upload", action: "other"}}, {method: http.MethodPost, path: "/v1/tokens", want: tenantRequestClass{surface: "tokens", action: "issue"}}, {method: http.MethodPost, path: "/v1/tokens/revoke", want: tenantRequestClass{surface: "tokens", action: "revoke_by_key"}}, {method: http.MethodDelete, path: "/v1/tokens/tok_123", want: tenantRequestClass{surface: "tokens", action: "revoke"}}, {method: http.MethodGet, path: "/v1/journals", want: tenantRequestClass{surface: "journal", action: "get"}}, {method: http.MethodPost, path: "/v1/journal-entries", want: tenantRequestClass{surface: "journal", action: "post"}}, {method: http.MethodGet, path: "/v1/git-workspaces/w1", want: tenantRequestClass{surface: "git_workspace", action: "get"}}, {method: http.MethodPost, path: "/v1/fork", want: tenantRequestClass{surface: "fork", action: "create"}}, {method: http.MethodDelete, path: "/v1/fork", want: tenantRequestClass{surface: "fork", action: "delete"}}, {method: http.MethodGet, path: "/v1/layers", want: tenantRequestClass{surface: "layer", action: "list"}}, {method: http.MethodPost, path: "/v1/layers", want: tenantRequestClass{surface: "layer", action: "create"}}, {method: http.MethodGet, path: "/v1/layers/l1", want: tenantRequestClass{surface: "layer", action: "get"}}, {method: http.MethodGet, path: "/v1/layers/l1/diff", want: tenantRequestClass{surface: "layer", action: "diff"}}, {method: http.MethodPost, path: "/v1/layers/l1/checkpoints", want: tenantRequestClass{surface: "layer", action: "checkpoint"}}, {method: http.MethodPost, path: "/v1/layers/l1/uploads/initiate", want: tenantRequestClass{surface: "layer", action: "upload_initiate"}}, {method: http.MethodPost, path: "/v1/layers/l1/uploads/u1/presign", want: tenantRequestClass{surface: "layer", action: "upload_presign"}}, {method: http.MethodPost, path: "/v1/layers/l1/uploads/u1/presign-batch", want: tenantRequestClass{surface: "layer", action: "upload_presign_batch"}}, {method: http.MethodPost, path: "/v1/layers/l1/uploads/u1/complete", want: tenantRequestClass{surface: "layer", action: "upload_complete"}}, {method: http.MethodPost, path: "/v1/layers/l1/uploads/u1/abort", want: tenantRequestClass{surface: "layer", action: "upload_abort"}}, {method: http.MethodGet, path: "/v1/layer-checkpoints/c1", want: tenantRequestClass{surface: "layer", action: "checkpoint_get"}}, {method: http.MethodPost, path: "/v1/admin/tenant-pool", want: tenantRequestClass{surface: "admin", action: "post"}}, {method: http.MethodGet, path: "/v1/admin/tenants/t1", want: tenantRequestClass{surface: "admin", action: "get"}}, {method: http.MethodGet, path: "/v1/auth/slock/login", want: tenantRequestClass{surface: "auth", action: "slock_login"}}, {method: http.MethodGet, path: "/v1/auth/slock/callback", want: tenantRequestClass{surface: "auth", action: "slock_callback"}}, {method: http.MethodGet, path: "/v1/vault/read/db/password", want: tenantRequestClass{surface: "vault", action: "read"}}, {method: http.MethodPut, path: "/s3/local/upload/u1/1", want: tenantRequestClass{surface: "object_store", action: "upload_part"}}, {method: http.MethodGet, path: "/s3/local/objects/blob", want: tenantRequestClass{surface: "object_store", action: "get_object"}}, {method: http.MethodPost, path: "/v1/sql", want: tenantRequestClass{surface: "sql", action: "post"}}, {method: http.MethodPost, path: "/v1/migration/events", want: tenantRequestClass{surface: "migration", action: "event"}}, } for _, tc := range tests { req, err := http.NewRequest(tc.method, "http://example.test"+tc.path, nil) if err != nil { t.Fatal(err) } got := classifyTenantRequest(req) if got != tc.want { t.Fatalf("classifyTenantRequest(%s %s) = %#v, want %#v", tc.method, tc.path, got, tc.want) } } } func TestRequestTenantID(t *testing.T) { t.Parallel() tests := []struct { path string want string }{ {path: "/s3/local/upload/u1/1", want: "local"}, {path: "/s3/upload/u1/1", want: "local"}, {path: "/s3/objects/blob", want: "local"}, {path: "/s3/tenant-a/upload/u1/1", want: ""}, {path: "/s3/tenant-a/objects/blob", want: ""}, {path: "/v1/fs/doc.txt", want: ""}, } for _, tc := range tests { req, err := http.NewRequest(http.MethodGet, "http://example.test"+tc.path, nil) if err != nil { t.Fatal(err) } if got := requestTenantID(req); got != tc.want { t.Fatalf("requestTenantID(%q) = %q, want %q", tc.path, got, tc.want) } } req, err := http.NewRequest(http.MethodGet, "http://example.test/s3/tenant-a/objects/blob", nil) if err != nil { t.Fatal(err) } ctx := withRequestMetricState(req.Context(), &requestMetricState{}) req = req.WithContext(ctx) setRequestMetricTenant(req.Context(), "tenant-a", "", "", "org-a", tenantRequestClass{surface: "object_store", action: "get_object"}) if got := requestTenantID(req); got != "tenant-a" { t.Fatalf("requestTenantID with verified scope = %q, want tenant-a", got) } } func TestSetRequestMetricTenantMovesInFlightLabel(t *testing.T) { m := newServerMetrics() state := &requestMetricState{} ctx := withMetrics(context.Background(), m) ctx = withRequestMetricState(ctx, state) class := tenantRequestClass{surface: "object_store", action: "upload_part"} setRequestMetricTenant(ctx, "local", "", "", "", class) if got := m.tenantInFlight[tenantInFlightKey("local", defaultTenantMetricTiDBCloudOrgID, "object_store")]; got != 1 { t.Fatalf("local in-flight = %d, want 1", got) } setRequestMetricTenant(ctx, "tenant-a", "", "", "org-a", class) if got := m.tenantInFlight[tenantInFlightKey("local", defaultTenantMetricTiDBCloudOrgID, "object_store")]; got != 0 { t.Fatalf("local in-flight after move = %d, want 0", got) } if got := m.tenantInFlight[tenantInFlightKey("tenant-a", "org-a", "object_store")]; got != 1 { t.Fatalf("tenant-a in-flight = %d, want 1", got) } finishRequestMetricTenant(ctx) if got := len(m.tenantInFlight); got != 0 { t.Fatalf("tenant in-flight map size after finish = %d, want 0", got) } } func TestSetRequestMetricTenantForAuthStatusOnlyScopesActiveTenants(t *testing.T) { for _, status := range []meta.TenantStatus{meta.TenantDeleting, meta.TenantDeleted, meta.TenantSuspended} { t.Run(string(status), func(t *testing.T) { const tenantID = "tenant-auth-status-metrics" metrics.DeleteTenantCounters(tenantID) t.Cleanup(func() { metrics.DeleteTenantCounters(tenantID) }) ctx := withRequestMetricState(context.Background(), &requestMetricState{}) setRequestMetricTenantForAuthStatus(ctx, tenantID, "key", "db9", "org", status, tenantRequestClass{surface: "api", action: "read"}) if tenantID, _, _, _ := requestMetricScope(ctx); tenantID != "" { t.Fatalf("tenant metric scope = %q for status %q, want empty", tenantID, status) } req := httptest.NewRequest(http.MethodGet, "http://example.test/v1/fs/stale", nil).WithContext(ctx) recordTenantHTTPRequest(req, http.StatusForbidden, 0, 0) recorder := httptest.NewRecorder() metrics.WritePrometheus(recorder) if strings.Contains(recorder.Body.String(), `tenant_id="`+tenantID+`"`) { t.Fatalf("rejected %s request recreated tenant metrics:\n%s", status, recorder.Body.String()) } }) } ctx := withRequestMetricState(context.Background(), &requestMetricState{}) setRequestMetricTenantForAuthStatus(ctx, "tenant-active", "key", "db9", "org", meta.TenantActive, tenantRequestClass{surface: "api", action: "read"}) if tenantID, _, _, _ := requestMetricScope(ctx); tenantID != "tenant-active" { t.Fatalf("tenant metric scope = %q, want tenant-active", tenantID) } } func TestRecordTenantHTTPRequestAtCompletionDoesNotRecreateCleanedCounter(t *testing.T) { const tenantID = "tenant-request-completion-cleanup" metrics.DeleteTenantCounters(tenantID) t.Cleanup(func() { metrics.DeleteTenantCounters(tenantID) }) ctx := withRequestMetricState(context.Background(), &requestMetricState{}) setRequestTenantHTTPRecorder(ctx, func() { metrics.RecordTenantRequestCountWithOrg(tenantID, "org-completion-cleanup", "fs", "read", http.StatusOK) }) // The auth middleware runs this before it releases the final pool reference. recordTenantHTTPRequestAtCompletion(ctx) metrics.DeleteTenantCounters(tenantID) // The outer observe fallback must not recreate a counter after explicit cleanup. recordTenantHTTPRequestAtCompletion(ctx) recorder := httptest.NewRecorder() metrics.WritePrometheus(recorder) if strings.Contains(recorder.Body.String(), `drive9_tenant_requests_total{`) && strings.Contains(recorder.Body.String(), `tenant_id="`+tenantID+`"`) { t.Fatalf("tenant request counter recreated after cleanup:\n%s", recorder.Body.String()) } } func TestAdjustTenantInFlightAggregatesActionsForSameSurface(t *testing.T) { m := newServerMetrics() if got := m.adjustTenantInFlight("tenant-inflight-actions", "org-inflight-actions", "fs", 1); got != 1 { t.Fatalf("read increment = %d, want 1", got) } if got := m.adjustTenantInFlight("tenant-inflight-actions", "org-inflight-actions", "fs", 1); got != 2 { t.Fatalf("write increment = %d, want aggregate 2", got) } if got := m.adjustTenantInFlight("tenant-inflight-actions", "org-inflight-actions", "fs", -1); got != 1 { t.Fatalf("read decrement = %d, want aggregate 1", got) } if got := m.adjustTenantInFlight("tenant-inflight-actions", "org-inflight-actions", "fs", -1); got != 0 { t.Fatalf("write decrement = %d, want 0", got) } } func TestEventFieldsDoesNotDuplicateExplicitTiDBCloudOrgID(t *testing.T) { ctx := withRequestMetricState(context.Background(), &requestMetricState{}) setRequestMetricTenant(ctx, "tenant-a", "", "", "org-request", tenantRequestClass{surface: "status", action: "get"}) fields := eventFields(ctx, "test_event", "tidbcloud_org_id", "org-explicit") count := 0 for _, field := range fields { if field.Key != "tidbcloud_org_id" { continue } count++ if field.String != "org-explicit" { t.Fatalf("tidbcloud_org_id = %q, want org-explicit", field.String) } } if count != 1 { t.Fatalf("tidbcloud_org_id field count = %d, want 1", count) } } func TestRequestTenantMetricScopeFallbackLeavesOrgEmpty(t *testing.T) { req, err := http.NewRequest(http.MethodGet, "http://example.test/s3/objects/blob", nil) if err != nil { t.Fatal(err) } req = req.WithContext(withRequestMetricState(req.Context(), &requestMetricState{})) tenantID, tidbCloudOrgID := requestTenantMetricScope(req) if tenantID != "local" { t.Fatalf("tenant id = %q, want local", tenantID) } if tidbCloudOrgID != "" { t.Fatalf("tidb cloud org id = %q, want empty", tidbCloudOrgID) } } type flushRecorder struct { header http.Header flushed bool } type plainRecorder struct { header http.Header } func (r *flushRecorder) Header() http.Header { if r.header == nil { r.header = make(http.Header) } return r.header } func (r *flushRecorder) Write(p []byte) (int, error) { return len(p), nil } func (r *flushRecorder) WriteHeader(_ int) {} func (r *flushRecorder) Flush() { r.flushed = true } func (r *plainRecorder) Header() http.Header { if r.header == nil { r.header = make(http.Header) } return r.header } func (r *plainRecorder) Write(p []byte) (int, error) { return len(p), nil } func (r *plainRecorder) WriteHeader(_ int) {} func TestObservedResponseWriterDoesNotAdvertiseFlush(t *testing.T) { ow := &observedResponseWriter{ResponseWriter: &plainRecorder{}} if _, ok := interface{}(ow).(http.Flusher); ok { t.Fatal("observedResponseWriter should not implement http.Flusher") } } func TestFlusherResponseWriterDelegatesFlush(t *testing.T) { rec := &flushRecorder{} fw := &flusherResponseWriter{ observedResponseWriter: &observedResponseWriter{ResponseWriter: rec}, flusher: rec, } if _, ok := interface{}(fw).(http.Flusher); !ok { t.Fatal("flusherResponseWriter should implement http.Flusher") } fw.Flush() if !rec.flushed { t.Fatal("expected Flush to delegate to wrapped writer") } } // TestSSEEstablishmentDuration verifies that markSSEStreamEstablished records // the establishment timestamp so observe can measure the time-to-first-byte // (request → 200 headers + flush) rather than the full SSE connection // lifetime. This is the core invariant that keeps /v1/events in the HTTP p99 // latency alert on a bounded, meaningful basis without the select-loop // hold-open time polluting the histogram. func TestSSEEstablishmentDuration(t *testing.T) { state := &requestMetricState{} ctx := withRequestMetricState(context.Background(), state) // Before establishment: flag is false and duration is unavailable. if sseStreamEstablished(ctx) { t.Fatal("sseStreamEstablished should be false before markSSEStreamEstablished") } start := time.Now() if _, ok := sseEstablishmentDuration(ctx, start); ok { t.Fatal("sseEstablishmentDuration should return false before establishment") } // Simulate the establishment point (handleEvents writing 200 headers). time.Sleep(2 * time.Millisecond) markSSEStreamEstablished(ctx) if !sseStreamEstablished(ctx) { t.Fatal("sseStreamEstablished should be true after markSSEStreamEstablished") } estDur, ok := sseEstablishmentDuration(ctx, start) if !ok { t.Fatal("sseEstablishmentDuration should return true after establishment") } if estDur < time.Millisecond { t.Fatalf("establishment duration = %v, want >= 1ms", estDur) } // The establishment duration must be bounded and significantly shorter // than a simulated connection lifetime. observe uses estDur (not // time.Since(start)) for the HTTP duration histogram when the stream is // established. connLifetime := time.Since(start) if estDur > connLifetime { t.Fatalf("establishment duration %v should not exceed connection lifetime %v", estDur, connLifetime) } }