72de690894
Chatbot functionality * baseline working * missing test file * more tests
595 lines
21 KiB
Go
595 lines
21 KiB
Go
// Package bot — Milestone 3 payload assembly tests.
|
||
//
|
||
// These tests cover the four payload-assembly helpers backend-eng adds
|
||
// in M3 (per plan §7 mapping table):
|
||
//
|
||
// - Conversation history: last RecentTurns completed turns × 2 messages,
|
||
// ascending ordinal.
|
||
// - Scope payload assembly: stored doc ids ∪ folder-expanded doc ids,
|
||
// with the original folder ids preserved in the payload.
|
||
// - Cached-results stripping on updated_session_state.
|
||
// - 64KB session-state cap (chatbot.StateMaxBytes).
|
||
// - Metadata summarization (sum input/output tokens across the
|
||
// metadata.agents map; wall-clock fallback when metadata is absent).
|
||
//
|
||
// SQL-driven helpers (conversation history, scope expansion) use the
|
||
// real Postgres testcontainer. Pure functions (strip, cap, summarize)
|
||
// run in-memory. No httptest.NewServer here — that lives in
|
||
// fastapi_test.go for transport-layer cases only.
|
||
//
|
||
// Compile contract: this file references bot.BuildConversationHistory,
|
||
// bot.BuildChatScope, bot.StripCachedResults, bot.SummarizeMetadata,
|
||
// bot.SessionStateOversize, and the previously-declared types from
|
||
// fastapi_test.go (ChatScope, ConversationMessage, Metadata,
|
||
// AgentTelemetry). Backend-eng adds payload.go / state_strip.go /
|
||
// metadata_summarise.go and the package compiles.
|
||
package bot_test
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"os"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"queryorchestration/internal/bot"
|
||
"queryorchestration/internal/database/repository"
|
||
"queryorchestration/internal/serviceconfig"
|
||
"queryorchestration/internal/serviceconfig/chatbot"
|
||
"queryorchestration/internal/test"
|
||
|
||
"github.com/google/uuid"
|
||
"github.com/stretchr/testify/assert"
|
||
"github.com/stretchr/testify/require"
|
||
)
|
||
|
||
// payloadDefaultCfg returns a chatbot.ChatbotConfig with the plan §8
|
||
// defaults so the test does not have to repeat the literal values. The
|
||
// SERVICE_URL/API_KEY values are placeholders; payload tests do not
|
||
// dial out.
|
||
func payloadDefaultCfg() chatbot.ChatbotConfig {
|
||
return chatbot.ChatbotConfig{
|
||
ServiceURL: "http://chatbot.test:8000",
|
||
APIKey: "test-api-key",
|
||
RequestTimeoutSeconds: 60,
|
||
MaxTurnChars: 2000,
|
||
RecentTurns: 5,
|
||
}
|
||
}
|
||
|
||
// resetClientForPayload mirrors resetClientForBotTests but lives in the
|
||
// bot package's test space so this file does not pull in queryapi_test
|
||
// helpers. The shared testcontainer Postgres reuses volumes across runs,
|
||
// so leftover bot rows must be wiped before CreateClient.
|
||
func resetClientForPayload(t *testing.T, ctx context.Context, cfg *serviceconfig.BaseConfig, clientID, name string) {
|
||
t.Helper()
|
||
pool := cfg.GetDBPool()
|
||
stmts := []string{
|
||
`DELETE FROM document_custom_metadata WHERE document_id IN (SELECT id FROM documents WHERE clientId = $1)`,
|
||
`UPDATE documents SET custom_schema_id = NULL WHERE clientId = $1`,
|
||
`DELETE FROM client_metadata_schemas WHERE client_id = $1`,
|
||
`DELETE FROM bot_session_documents WHERE document_id IN (SELECT id FROM documents WHERE clientId = $1)`,
|
||
`DELETE FROM bot_session_folders WHERE folder_id IN (SELECT id FROM folders WHERE clientId = $1)`,
|
||
`DELETE FROM bot_sessions WHERE client_id = $1`,
|
||
`DELETE FROM documents WHERE clientId = $1`,
|
||
`DELETE FROM folders WHERE clientId = $1`,
|
||
`DELETE FROM clients WHERE clientId = $1`,
|
||
}
|
||
for _, stmt := range stmts {
|
||
_, err := pool.Exec(ctx, stmt, clientID)
|
||
require.NoErrorf(t, err, "reset stmt failed: %s", stmt)
|
||
}
|
||
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
||
Name: name,
|
||
Clientid: clientID,
|
||
})
|
||
require.NoError(t, err)
|
||
}
|
||
|
||
// seedPayloadSession inserts a bot_sessions row owned by (clientID, actor).
|
||
func seedPayloadSession(t *testing.T, ctx context.Context, cfg *serviceconfig.BaseConfig, clientID, actor string) uuid.UUID {
|
||
t.Helper()
|
||
var id uuid.UUID
|
||
err := cfg.GetDBPool().QueryRow(ctx, `
|
||
INSERT INTO bot_sessions (client_id, created_by) VALUES ($1, $2) RETURNING id
|
||
`, clientID, actor).Scan(&id)
|
||
require.NoError(t, err)
|
||
return id
|
||
}
|
||
|
||
// seedCompletedTurn inserts one completed bot_turns row at the given
|
||
// ordinal, with prompt and completion populated.
|
||
func seedCompletedTurn(t *testing.T, ctx context.Context, cfg *serviceconfig.BaseConfig, sessionID uuid.UUID, ordinal int, prompt, completion string, createdAt time.Time) {
|
||
t.Helper()
|
||
_, err := cfg.GetDBPool().Exec(ctx, `
|
||
INSERT INTO bot_turns (session_id, ordinal, prompt, status, attempt_id, completion, attempt_started_at, created_at)
|
||
VALUES ($1, $2, $3, 'completed', $4, $5, $6, $6)
|
||
`, sessionID, ordinal, prompt, uuid.New(), completion, createdAt)
|
||
require.NoError(t, err)
|
||
}
|
||
|
||
// ---------- conversation history ----------
|
||
|
||
// TestConversationHistoryConversion: given the last RecentTurns
|
||
// completed turns, BuildConversationHistory produces N×2 messages
|
||
// (user prompt + assistant completion) in ascending ordinal order.
|
||
// Plan §7 mapping table.
|
||
func TestConversationHistoryConversion(t *testing.T) {
|
||
if testing.Short() {
|
||
t.SkipNow()
|
||
}
|
||
ctx := t.Context()
|
||
|
||
dbCfg := &serviceconfig.BaseConfig{}
|
||
test.CreateDB(t, dbCfg)
|
||
|
||
const (
|
||
clientID = "BOT_M3_CONV"
|
||
actor = "history@example.com"
|
||
)
|
||
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-conv")
|
||
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
|
||
|
||
// Seed 7 completed turns. With cfg.RecentTurns=5 the helper must
|
||
// pick the last 5 ordered ascending: ordinals 3..7. createdAt is
|
||
// staggered by one second so any future ordering by created_at
|
||
// (rather than ordinal) is also deterministic.
|
||
now := time.Now().UTC().Truncate(time.Second)
|
||
for ord := 1; ord <= 7; ord++ {
|
||
seedCompletedTurn(t, ctx, dbCfg, sessionID, ord,
|
||
"prompt-"+string(rune('0'+ord)),
|
||
"answer-"+string(rune('0'+ord)),
|
||
now.Add(time.Duration(ord)*time.Second))
|
||
}
|
||
|
||
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
|
||
require.NoError(t, err)
|
||
|
||
history, err := svc.BuildConversationHistory(ctx, sessionID)
|
||
require.NoError(t, err)
|
||
require.Len(t, history, 5*2,
|
||
"plan §7: RecentTurns=5 -> 5 turns x 2 messages = 10")
|
||
|
||
// Expect ordinals 3,4,5,6,7 in ascending order. Each turn produces
|
||
// (user prompt, assistant completion) in that order.
|
||
for i := 0; i < 5; i++ {
|
||
ord := i + 3
|
||
userMsg := history[i*2]
|
||
assistantMsg := history[i*2+1]
|
||
|
||
assert.Equal(t, "user", userMsg.Role,
|
||
"index %d must be user; ordinal=%d", i*2, ord)
|
||
assert.Equal(t, "prompt-"+string(rune('0'+ord)), userMsg.Content,
|
||
"user content for ordinal %d", ord)
|
||
assert.Equal(t, "assistant", assistantMsg.Role,
|
||
"index %d must be assistant; ordinal=%d", i*2+1, ord)
|
||
assert.Equal(t, "answer-"+string(rune('0'+ord)), assistantMsg.Content,
|
||
"assistant content for ordinal %d", ord)
|
||
assert.Falsef(t, userMsg.Timestamp.IsZero(),
|
||
"user message ordinal %d must carry a non-zero timestamp", ord)
|
||
}
|
||
}
|
||
|
||
// TestConversationHistoryConversion_EmptySession: a session with zero
|
||
// completed turns yields a zero-length slice (NOT nil), so the FastAPI
|
||
// client emits the JSON `"conversation_history": []` per the swagger
|
||
// requirement that the field be a non-null array.
|
||
func TestConversationHistoryConversion_EmptySession(t *testing.T) {
|
||
if testing.Short() {
|
||
t.SkipNow()
|
||
}
|
||
ctx := t.Context()
|
||
|
||
dbCfg := &serviceconfig.BaseConfig{}
|
||
test.CreateDB(t, dbCfg)
|
||
|
||
const (
|
||
clientID = "BOT_M3_CONV_E"
|
||
actor = "empty@example.com"
|
||
)
|
||
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-conv-e")
|
||
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
|
||
|
||
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
|
||
require.NoError(t, err)
|
||
|
||
history, err := svc.BuildConversationHistory(ctx, sessionID)
|
||
require.NoError(t, err)
|
||
require.NotNil(t, history,
|
||
"empty result must be a non-nil slice so json.Marshal emits []")
|
||
assert.Empty(t, history)
|
||
}
|
||
|
||
// ---------- scope payload assembly ----------
|
||
|
||
// TestScopePayloadAssembly_FoldersExpanded: stored scope has 2 doc ids
|
||
// + 1 folder id; folder contains 3 unrelated docs. Resulting payload
|
||
// has documents = the union of 5 ids; folders = the original 1 folder.
|
||
// Plan §7 mapping table: "scope": "Stored document ids union expanded
|
||
// folder document ids; stored folder ids preserved".
|
||
func TestScopePayloadAssembly_FoldersExpanded(t *testing.T) {
|
||
if testing.Short() {
|
||
t.SkipNow()
|
||
}
|
||
ctx := t.Context()
|
||
|
||
dbCfg := &serviceconfig.BaseConfig{}
|
||
test.CreateDB(t, dbCfg)
|
||
pool := dbCfg.GetDBPool()
|
||
|
||
const (
|
||
clientID = "BOT_M3_SCOPE_FX"
|
||
actor = "scope@example.com"
|
||
)
|
||
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-scope-fx")
|
||
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
|
||
|
||
queries := dbCfg.GetDBQueries()
|
||
|
||
// 2 stored documents (no folder).
|
||
docA, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{Clientid: clientID, Hash: "scope-fx-a"})
|
||
require.NoError(t, err)
|
||
docB, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{Clientid: clientID, Hash: "scope-fx-b"})
|
||
require.NoError(t, err)
|
||
|
||
// 1 folder containing 3 documents (different from the stored docs).
|
||
folder, err := queries.CreateFolder(ctx, &repository.CreateFolderParams{
|
||
Path: "/scope-fx-folder", Parentid: nil, Clientid: clientID, Createdby: "tester",
|
||
})
|
||
require.NoError(t, err)
|
||
folderDocs := make([]uuid.UUID, 3)
|
||
for i := 0; i < 3; i++ {
|
||
var id uuid.UUID
|
||
err := pool.QueryRow(ctx, `
|
||
INSERT INTO documents (clientId, hash, folderId) VALUES ($1, $2, $3) RETURNING id
|
||
`, clientID, "scope-fx-folderdoc-"+string(rune('0'+i)), folder.ID).Scan(&id)
|
||
require.NoError(t, err)
|
||
folderDocs[i] = id
|
||
}
|
||
|
||
// Persist the scope: 2 stored docs + 1 folder.
|
||
for _, d := range []uuid.UUID{docA, docB} {
|
||
_, err := pool.Exec(ctx,
|
||
`INSERT INTO bot_session_documents (session_id, document_id) VALUES ($1, $2)`,
|
||
sessionID, d)
|
||
require.NoError(t, err)
|
||
}
|
||
_, err = pool.Exec(ctx,
|
||
`INSERT INTO bot_session_folders (session_id, folder_id) VALUES ($1, $2)`,
|
||
sessionID, folder.ID)
|
||
require.NoError(t, err)
|
||
|
||
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
|
||
require.NoError(t, err)
|
||
|
||
scope, err := svc.BuildChatScope(ctx, sessionID, clientID)
|
||
require.NoError(t, err)
|
||
|
||
wantDocs := []string{
|
||
docA.String(), docB.String(),
|
||
folderDocs[0].String(), folderDocs[1].String(), folderDocs[2].String(),
|
||
}
|
||
assert.ElementsMatch(t, wantDocs, scope.DocumentIDs,
|
||
"plan §7: documents = stored docs ∪ folder-expanded docs")
|
||
assert.ElementsMatch(t, []string{folder.ID.String()}, scope.FolderIDs,
|
||
"plan §7: stored folder ids preserved")
|
||
}
|
||
|
||
// TestScopePayloadAssembly_DocsAndFolderOverlap: stored doc id is also
|
||
// inside the folder being expanded. Payload documents must dedup so the
|
||
// shared id appears exactly once. Folder id preserved.
|
||
func TestScopePayloadAssembly_DocsAndFolderOverlap(t *testing.T) {
|
||
if testing.Short() {
|
||
t.SkipNow()
|
||
}
|
||
ctx := t.Context()
|
||
|
||
dbCfg := &serviceconfig.BaseConfig{}
|
||
test.CreateDB(t, dbCfg)
|
||
pool := dbCfg.GetDBPool()
|
||
|
||
const (
|
||
clientID = "BOT_M3_SCOPE_OL"
|
||
actor = "overlap@example.com"
|
||
)
|
||
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-scope-ol")
|
||
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
|
||
|
||
queries := dbCfg.GetDBQueries()
|
||
|
||
folder, err := queries.CreateFolder(ctx, &repository.CreateFolderParams{
|
||
Path: "/scope-ol-folder", Parentid: nil, Clientid: clientID, Createdby: "tester",
|
||
})
|
||
require.NoError(t, err)
|
||
|
||
// docA is the overlap: stored AND inside the folder.
|
||
var docA uuid.UUID
|
||
err = pool.QueryRow(ctx, `
|
||
INSERT INTO documents (clientId, hash, folderId) VALUES ($1, $2, $3) RETURNING id
|
||
`, clientID, "scope-ol-overlap", folder.ID).Scan(&docA)
|
||
require.NoError(t, err)
|
||
|
||
// docB is folder-only.
|
||
var docB uuid.UUID
|
||
err = pool.QueryRow(ctx, `
|
||
INSERT INTO documents (clientId, hash, folderId) VALUES ($1, $2, $3) RETURNING id
|
||
`, clientID, "scope-ol-folder-only", folder.ID).Scan(&docB)
|
||
require.NoError(t, err)
|
||
|
||
// Stored scope: docA (also in folder) + folder.
|
||
_, err = pool.Exec(ctx,
|
||
`INSERT INTO bot_session_documents (session_id, document_id) VALUES ($1, $2)`,
|
||
sessionID, docA)
|
||
require.NoError(t, err)
|
||
_, err = pool.Exec(ctx,
|
||
`INSERT INTO bot_session_folders (session_id, folder_id) VALUES ($1, $2)`,
|
||
sessionID, folder.ID)
|
||
require.NoError(t, err)
|
||
|
||
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
|
||
require.NoError(t, err)
|
||
|
||
scope, err := svc.BuildChatScope(ctx, sessionID, clientID)
|
||
require.NoError(t, err)
|
||
|
||
wantDocs := []string{docA.String(), docB.String()}
|
||
assert.ElementsMatch(t, wantDocs, scope.DocumentIDs,
|
||
"overlap document must appear exactly once after dedup")
|
||
assert.ElementsMatch(t, []string{folder.ID.String()}, scope.FolderIDs)
|
||
}
|
||
|
||
// TestScopePayloadAssembly_EmptyScopeReturnsAllClientDocs: with no
|
||
// stored doc/folder rows, the payload documents fall back to "all
|
||
// documents for the client" per plan §4 ("Empty scope tables mean all
|
||
// documents for the client"). Folders array is empty.
|
||
func TestScopePayloadAssembly_EmptyScopeReturnsAllClientDocs(t *testing.T) {
|
||
if testing.Short() {
|
||
t.SkipNow()
|
||
}
|
||
ctx := t.Context()
|
||
|
||
dbCfg := &serviceconfig.BaseConfig{}
|
||
test.CreateDB(t, dbCfg)
|
||
|
||
const (
|
||
clientID = "BOT_M3_SCOPE_E"
|
||
actor = "empty-scope@example.com"
|
||
)
|
||
resetClientForPayload(t, ctx, dbCfg, clientID, "bot-m3-scope-e")
|
||
sessionID := seedPayloadSession(t, ctx, dbCfg, clientID, actor)
|
||
|
||
queries := dbCfg.GetDBQueries()
|
||
wantDocs := []string{}
|
||
for i := 0; i < 4; i++ {
|
||
id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
|
||
Clientid: clientID, Hash: "scope-e-" + string(rune('0'+i)),
|
||
})
|
||
require.NoError(t, err)
|
||
wantDocs = append(wantDocs, id.String())
|
||
}
|
||
|
||
svc, err := bot.New(payloadDefaultCfg(), dbCfg)
|
||
require.NoError(t, err)
|
||
|
||
scope, err := svc.BuildChatScope(ctx, sessionID, clientID)
|
||
require.NoError(t, err)
|
||
assert.ElementsMatch(t, wantDocs, scope.DocumentIDs,
|
||
"empty scope tables -> all documents for the client")
|
||
assert.Empty(t, scope.FolderIDs)
|
||
}
|
||
|
||
// ---------- cached_results stripping ----------
|
||
|
||
// TestStripCachedResults_RemovesKey: input state with a cached_results
|
||
// entry returns the same object minus that one key. Other keys are
|
||
// preserved verbatim.
|
||
func TestStripCachedResults_RemovesKey(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
in := json.RawMessage(`{"foo":1,"cached_results":{"abc":{"k":"v"}},"bar":"hello"}`)
|
||
out, err := bot.StripCachedResults(in)
|
||
require.NoError(t, err)
|
||
|
||
var parsed map[string]interface{}
|
||
require.NoError(t, json.Unmarshal(out, &parsed))
|
||
_, hasCached := parsed["cached_results"]
|
||
assert.False(t, hasCached, "cached_results must be removed")
|
||
assert.EqualValues(t, 1, parsed["foo"], "other keys must be preserved")
|
||
assert.Equal(t, "hello", parsed["bar"])
|
||
}
|
||
|
||
// TestStripCachedResults_NoOpWhenAbsent: input without cached_results is
|
||
// returned with the same logical shape.
|
||
func TestStripCachedResults_NoOpWhenAbsent(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
in := json.RawMessage(`{"foo":1,"bar":"hello"}`)
|
||
out, err := bot.StripCachedResults(in)
|
||
require.NoError(t, err)
|
||
|
||
var got map[string]interface{}
|
||
require.NoError(t, json.Unmarshal(out, &got))
|
||
|
||
var want map[string]interface{}
|
||
require.NoError(t, json.Unmarshal(in, &want))
|
||
|
||
assert.Equal(t, want, got)
|
||
}
|
||
|
||
// TestStripCachedResults_NullInput: a nil/empty input passes through
|
||
// unchanged so the caller can hand it null state without a special case.
|
||
func TestStripCachedResults_NullInput(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
out, err := bot.StripCachedResults(nil)
|
||
require.NoError(t, err)
|
||
assert.Nil(t, out, "nil input must round-trip as nil")
|
||
}
|
||
|
||
// ---------- 64KB session-state cap ----------
|
||
|
||
// TestSessionStateSizeCap_64KB_Boundary: the cap is plan §8's
|
||
// chatbot.StateMaxBytes. A stripped state at exactly the cap reports
|
||
// as not-oversize; one byte over reports oversize. The helper returns
|
||
// (oversize bool) without mutating the input.
|
||
func TestSessionStateSizeCap_64KB_Boundary(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
// Build a JSON object whose length is exactly chatbot.StateMaxBytes.
|
||
// We use a single-key string-valued object and pad the value so the
|
||
// total raw byte length is exact.
|
||
const head = `{"k":"`
|
||
const tail = `"}`
|
||
overhead := len(head) + len(tail)
|
||
require.Greater(t, chatbot.StateMaxBytes, overhead)
|
||
value := strings.Repeat("a", chatbot.StateMaxBytes-overhead)
|
||
stateAtCap := json.RawMessage(head + value + tail)
|
||
require.Equal(t, chatbot.StateMaxBytes, len(stateAtCap),
|
||
"test setup: stateAtCap must equal the cap byte-for-byte")
|
||
|
||
assert.False(t, bot.SessionStateOversize(stateAtCap),
|
||
"plan §8: state at exactly StateMaxBytes is not oversize")
|
||
|
||
// One byte over the cap.
|
||
overValue := strings.Repeat("a", chatbot.StateMaxBytes-overhead+1)
|
||
stateOver := json.RawMessage(head + overValue + tail)
|
||
require.Equal(t, chatbot.StateMaxBytes+1, len(stateOver))
|
||
|
||
assert.True(t, bot.SessionStateOversize(stateOver),
|
||
"plan §8: state above StateMaxBytes is oversize")
|
||
}
|
||
|
||
// ---------- metadata summarization ----------
|
||
|
||
// TestMetadataSummarization: response metadata.agents = {a: {input:10,
|
||
// output:20}, b: {input:5, output:30}} -> tokens_in=15, tokens_out=50,
|
||
// latency = metadata.total_latency_ms (NOT wall clock when metadata is
|
||
// present).
|
||
func TestMetadataSummarization(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
sentAt := time.Date(2026, 5, 4, 12, 0, 0, 0, time.UTC)
|
||
receivedAt := sentAt.Add(7 * time.Second)
|
||
|
||
metadata := &bot.Metadata{
|
||
TotalLatencyMs: 6000,
|
||
RetryCount: 0,
|
||
Agents: map[string]bot.AgentTelemetry{
|
||
"a": {LatencyMs: 1000, InputTokens: 10, OutputTokens: 20},
|
||
"b": {LatencyMs: 5000, InputTokens: 5, OutputTokens: 30},
|
||
},
|
||
}
|
||
|
||
latencyMs, tokensIn, tokensOut := bot.SummarizeMetadata(metadata, sentAt, receivedAt)
|
||
assert.Equal(t, 6000, latencyMs,
|
||
"plan §7: latency = metadata.total_latency_ms when metadata is present")
|
||
assert.Equal(t, 15, tokensIn,
|
||
"plan §7: tokens_in = sum of agents[].input_tokens")
|
||
assert.Equal(t, 50, tokensOut,
|
||
"plan §7: tokens_out = sum of agents[].output_tokens")
|
||
}
|
||
|
||
// TestMetadataSummarization_NilMetadataUsesWallClock: nil metadata ->
|
||
// latency = ReceivedAt - SentAt in ms, tokens default to 0.
|
||
func TestMetadataSummarization_NilMetadataUsesWallClock(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
sentAt := time.Date(2026, 5, 4, 12, 0, 0, 0, time.UTC)
|
||
receivedAt := sentAt.Add(2500 * time.Millisecond)
|
||
|
||
latencyMs, tokensIn, tokensOut := bot.SummarizeMetadata(nil, sentAt, receivedAt)
|
||
assert.Equal(t, 2500, latencyMs,
|
||
"nil metadata -> wall-clock latency in milliseconds")
|
||
assert.Equal(t, 0, tokensIn)
|
||
assert.Equal(t, 0, tokensOut)
|
||
}
|
||
|
||
// TestMetadataSummarization_EmptyAgentsMap: metadata present but agents
|
||
// map is empty -> tokens default to 0; latency uses metadata.total_latency_ms.
|
||
func TestMetadataSummarization_EmptyAgentsMap(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
sentAt := time.Date(2026, 5, 4, 12, 0, 0, 0, time.UTC)
|
||
receivedAt := sentAt.Add(5 * time.Second)
|
||
|
||
metadata := &bot.Metadata{
|
||
TotalLatencyMs: 4500,
|
||
RetryCount: 2,
|
||
Agents: map[string]bot.AgentTelemetry{},
|
||
}
|
||
|
||
latencyMs, tokensIn, tokensOut := bot.SummarizeMetadata(metadata, sentAt, receivedAt)
|
||
assert.Equal(t, 4500, latencyMs)
|
||
assert.Equal(t, 0, tokensIn)
|
||
assert.Equal(t, 0, tokensOut)
|
||
}
|
||
|
||
// TestStripCachedResults_OnRealAariaSample exercises StripCachedResults
|
||
// against the authoritative Aaria session_state shape Q committed to
|
||
// plans/chatbot.stuff/sample.response.json. The sample carries the full
|
||
// 10-key session-state envelope including a populated cached_results
|
||
// map keyed by result hash; stripping must leave the other 9 keys
|
||
// intact and the resulting bytes must fit comfortably under
|
||
// chatbot.StateMaxBytes.
|
||
//
|
||
// The other-key preservation check uses json.Unmarshal->reflect.DeepEqual
|
||
// rather than byte equality because Go's encoding/json does not
|
||
// guarantee map-key emission order: a sound implementation may re-emit
|
||
// keys in a different order than the source file, but the parsed value
|
||
// must be identical. Plan §7 mapping table, plan §8 64KB cap.
|
||
func TestStripCachedResults_OnRealAariaSample(t *testing.T) {
|
||
t.Parallel()
|
||
|
||
const samplePath = "../../plans/chatbot.stuff/sample.response.json"
|
||
raw, err := os.ReadFile(samplePath)
|
||
require.NoErrorf(t, err, "must read fixture at %s", samplePath)
|
||
|
||
// Sanity-precondition: the fixture has the 10 documented top-level
|
||
// keys including cached_results. If Q reshapes the fixture, this
|
||
// assertion fails first with a clear pointer at the source of drift.
|
||
var input map[string]interface{}
|
||
require.NoError(t, json.Unmarshal(raw, &input),
|
||
"fixture must be a JSON object")
|
||
require.Len(t, input, 10,
|
||
"fixture must have 10 top-level keys before stripping; got %d", len(input))
|
||
_, hadCached := input["cached_results"]
|
||
require.True(t, hadCached, "fixture must include cached_results")
|
||
|
||
stripped, err := bot.StripCachedResults(json.RawMessage(raw))
|
||
require.NoError(t, err)
|
||
require.NotNil(t, stripped, "non-empty input must produce non-nil output")
|
||
|
||
var out map[string]interface{}
|
||
require.NoError(t, json.Unmarshal(stripped, &out),
|
||
"stripped output must be a valid JSON object")
|
||
|
||
require.Len(t, out, 9,
|
||
"stripping cached_results must leave exactly 9 keys; got %d", len(out))
|
||
_, hasCached := out["cached_results"]
|
||
assert.False(t, hasCached, "cached_results must be absent after strip")
|
||
|
||
// Every other key must round-trip with the original value. Plan §7
|
||
// is explicit that ONLY cached_results is touched.
|
||
expectedKeys := []string{
|
||
"turn",
|
||
"resolved_entities",
|
||
"last_scope",
|
||
"last_intent",
|
||
"last_query_sql",
|
||
"last_query_plan",
|
||
"last_result_hash",
|
||
"last_scope_signature",
|
||
"validator_retry_count",
|
||
}
|
||
for _, k := range expectedKeys {
|
||
require.Containsf(t, out, k,
|
||
"key %q must be preserved after strip", k)
|
||
assert.Equalf(t, input[k], out[k],
|
||
"key %q value must round-trip unchanged", k)
|
||
}
|
||
|
||
assert.Falsef(t, bot.SessionStateOversize(stripped),
|
||
"plan §8: real Aaria sample is well under 64 KiB after strip; got %d bytes", len(stripped))
|
||
}
|