refactor: add dbgen chat generators and migrate test boilerplate (#24497)

- Adds chat-related dbgen generators covering defaults, overrides, and message field mapping.
- Replaces raw single-row chat, message, provider, and model-config setup in tests with dbgen helpers.
- Simplifies chat seed helpers after moving fixture setup into dbgen.

> Generated with [Coder Agents](https://coder.com/agents).

enterprise/coderd/x/chatd/chatd_retry_test.go

@@ -126,8 +126,8 @@ func TestRelayReconnectUsesExponentialBackoff(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-backoff")
+	user, org, model := seedChatDependencies(t, db)
+	chat := seedWaitingChat(t, db, org.ID, user, model, "relay-backoff")
 
 	_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
 	require.True(t, ok)
@@ -214,8 +214,8 @@ func TestRelayRepeatedDropsHitCap(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-drops")
+	user, org, model := seedChatDependencies(t, db)
+	chat := seedWaitingChat(t, db, org.ID, user, model, "relay-drops")
 
 	_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
 	require.True(t, ok)
@@ -305,8 +305,8 @@ func TestRelayStopsAfterIntermittentCap(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-cap")
+	user, org, model := seedChatDependencies(t, db)
+	chat := seedWaitingChat(t, db, org.ID, user, model, "relay-cap")
 
 	_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
 	require.True(t, ok)
@@ -416,8 +416,8 @@ func TestRelayReconnectStopsAfterDBErrorCap(t *testing.T) {
 	failingDB.okRemain.Store(1)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, realDB)
-	chat := seedWaitingChat(ctx, t, realDB, org.ID, user, model, "relay-db-error")
+	user, org, model := seedChatDependencies(t, realDB)
+	chat := seedWaitingChat(t, realDB, org.ID, user, model, "relay-db-error")
 
 	subscriber := newTestServer(t, failingDB, ps, subscriberID, dialer, mclk)
 	_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
@@ -509,8 +509,8 @@ func TestRelayStopsImmediatelyOnUnauthorized(t *testing.T) {
 			subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk)
 
 			ctx := testutil.Context(t, testutil.WaitLong)
-			user, org, model := seedChatDependencies(ctx, t, db)
-			chat := seedWaitingChat(ctx, t, db, org.ID, user, model,
+			user, org, model := seedChatDependencies(t, db)
+			chat := seedWaitingChat(t, db, org.ID, user, model,
 				"relay-unrec-"+tc.name)
 
 			_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
@@ -584,8 +584,8 @@ func TestRelayBackoffResetsOnStatusChange(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-reset-on-status")
+	user, org, model := seedChatDependencies(t, db)
+	chat := seedWaitingChat(t, db, org.ID, user, model, "relay-reset-on-status")
 
 	_, _, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
 	require.True(t, ok)
@@ -661,8 +661,8 @@ func TestRelayBackoffRespectsContextCancel(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-cancel")
+	user, org, model := seedChatDependencies(t, db)
+	chat := seedWaitingChat(t, db, org.ID, user, model, "relay-cancel")
 
 	subCtx, subCancel := context.WithCancel(ctx)
 	_, events, cancel, ok := subscriber.Subscribe(subCtx, chat.ID, nil, 0)
@@ -740,11 +740,11 @@ func TestDialRelayReal401(t *testing.T) {
 	subscribeFn := entchatd.NewMultiReplicaSubscribeFn(cfg)
 
 	ctx := testutil.Context(t, testutil.WaitMedium)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 	// Seed a waiting chat - no sync dial - then push a running
 	// status notification to trigger the async dial via the real
 	// dialRelay path.
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-real-401")
+	chat := seedWaitingChat(t, db, org.ID, user, model, "relay-real-401")
 
 	statusCh := make(chan osschatd.StatusNotification, 1)
 	evs := subscribeFn(ctx, osschatd.SubscribeFnParams{

enterprise/coderd/x/chatd/chatd_test.go

@@ -91,7 +91,6 @@ func newActiveWorkerServer(
 // seedChatDependencies creates a user, organization, and chat model
 // config in the database for use in relay tests.
 func seedChatDependencies(
-	ctx context.Context,
 	t *testing.T,
 	db database.Store,
 ) (database.User, database.Organization, database.ChatModelConfig) {
@@ -110,35 +109,19 @@ func seedChatDependencies(
 		UserID:         user.ID,
 		OrganizationID: org.ID,
 	})
-	_, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{
-		Provider:             "openai",
-		DisplayName:          "OpenAI",
-		APIKey:               "test-key",
-		BaseUrl:              safetyNet.URL,
-		CentralApiKeyEnabled: true,
-		ApiKeyKeyID:          sql.NullString{},
-		CreatedBy:            uuid.NullUUID{UUID: user.ID, Valid: true},
-		Enabled:              true,
+	_ = dbgen.ChatProvider(t, db, database.ChatProvider{
+		BaseUrl:   safetyNet.URL,
+		CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
 	})
-	require.NoError(t, err)
-	model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
-		Provider:             "openai",
-		Model:                "gpt-4o-mini",
-		DisplayName:          "Test Model",
-		CreatedBy:            uuid.NullUUID{UUID: user.ID, Valid: true},
-		UpdatedBy:            uuid.NullUUID{UUID: user.ID, Valid: true},
-		Enabled:              true,
-		IsDefault:            true,
-		ContextLimit:         128000,
-		CompressionThreshold: 70,
-		Options:              json.RawMessage(`{}`),
+	model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
+		CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
+		UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
+		IsDefault: true,
 	})
-	require.NoError(t, err)
 	return user, org, model
 }
 
 func seedWaitingChat(
-	ctx context.Context,
 	t *testing.T,
 	db database.Store,
 	orgID uuid.UUID,
@@ -148,16 +131,12 @@ func seedWaitingChat(
 ) database.Chat {
 	t.Helper()
 
-	chat, err := db.InsertChat(ctx, database.InsertChatParams{
+	chat := dbgen.Chat(t, db, database.Chat{
 		OrganizationID:    orgID,
-		Status:            database.ChatStatusWaiting,
-		ClientType:        database.ChatClientTypeUi,
 		OwnerID:           user.ID,
 		LastModelConfigID: model.ID,
 		Title:             title,
-		MCPServerIDs:      []uuid.UUID{},
 	})
-	require.NoError(t, err)
 	return chat
 }
 
@@ -173,7 +152,7 @@ func seedRemoteRunningChat(
 ) database.Chat {
 	t.Helper()
 
-	chat := seedWaitingChat(ctx, t, db, orgID, user, model, title)
+	chat := seedWaitingChat(t, db, orgID, user, model, title)
 	now := time.Now()
 	chat, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
 		ID:          chat.ID,
@@ -258,7 +237,7 @@ func TestSubscribeRelayReconnectsOnDrop(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 
 	chat := seedRemoteRunningChat(ctx, t, db, org.ID, user, model, workerID, "relay-reconnect")
 
@@ -336,11 +315,11 @@ func TestSubscribeRelayAsyncDoesNotBlock(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, provider, nil)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 
 	// Seed a waiting chat so Subscribe does not trigger a synchronous
 	// relay.
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-async-nonblock")
+	chat := seedWaitingChat(t, db, org.ID, user, model, "relay-async-nonblock")
 
 	// Subscribe before the chat is marked running so the relay opens
 	// via pubsub notification (openRelayAsync path).
@@ -438,7 +417,7 @@ func TestSubscribeRelaySnapshotDelivered(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, provider, nil)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 
 	chat := seedRemoteRunningChat(ctx, t, db, org.ID, user, model, workerID, "relay-snapshot")
 
@@ -526,7 +505,7 @@ func TestSubscribeRetryEventAcrossInstances(t *testing.T) {
 	}, nil)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 	setOpenAIProviderBaseURL(ctx, t, db, openAIURL)
 
 	chat, err := worker.CreateChat(ctx, osschatd.CreateOptions{
@@ -663,11 +642,11 @@ func TestSubscribeRelayStaleDialDiscardedAfterInterrupt(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, provider, nil)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 
 	// Seed the chat in waiting state so Subscribe does not try an initial
 	// relay.
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "stale-dial-test")
+	chat := seedWaitingChat(t, db, org.ID, user, model, "stale-dial-test")
 
 	// Subscribe while chat is in "waiting" state — no relay opened.
 	_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
@@ -815,11 +794,11 @@ func TestSubscribeCancelDuringInFlightDial(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, provider, nil)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 
 	// Seed the chat in waiting state so Subscribe does not open a
 	// synchronous relay.
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "cancel-inflight-dial")
+	chat := seedWaitingChat(t, db, org.ID, user, model, "cancel-inflight-dial")
 
 	_, _, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
 	require.True(t, ok)
@@ -901,10 +880,10 @@ func TestSubscribeRelayRunningToRunningSwitch(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, provider, nil)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 
 	// Seed the chat in waiting state so Subscribe does not open a relay.
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "running-to-running")
+	chat := seedWaitingChat(t, db, org.ID, user, model, "running-to-running")
 
 	_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
 	require.True(t, ok)
@@ -1009,11 +988,11 @@ func TestSubscribeRelayFailedDialRetries(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 
 	// Seed the chat in waiting state so Subscribe does not open a
 	// synchronous relay dial.
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "failed-dial-retry")
+	chat := seedWaitingChat(t, db, org.ID, user, model, "failed-dial-retry")
 
 	_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
 	require.True(t, ok)
@@ -1105,7 +1084,7 @@ func TestSubscribeRunningLocalWorkerClosesRelay(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, provider, nil)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 
 	chat := seedRemoteRunningChat(
 		ctx,
@@ -1205,7 +1184,7 @@ func TestSubscribeRelayMultipleReconnects(t *testing.T) {
 	subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 
 	chat := seedRemoteRunningChat(
 		ctx,
@@ -1349,13 +1328,13 @@ func TestSubscribeRelayDialCanceledOnFastCompletion(t *testing.T) {
 	}, nil)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 	setOpenAIProviderBaseURL(ctx, t, db, openAIURL)
 
 	// Create the chat in waiting state so the subscriber sees it
 	// before the worker picks it up (avoids the synchronous relay
 	// path in Subscribe).
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "fast-completion-relay-race")
+	chat := seedWaitingChat(t, db, org.ID, user, model, "fast-completion-relay-race")
 
 	// Subscribe from the subscriber replica while the chat is idle.
 	// No relay is opened because the chat is in waiting state.
@@ -1505,10 +1484,10 @@ func TestSubscribeRelayDrainWithinGraceLeavesBufferRetained(t *testing.T) {
 	}, subscriberClock)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 	setOpenAIProviderBaseURL(ctx, t, db, openAIURL)
 
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-drain-characterization")
+	chat := seedWaitingChat(t, db, org.ID, user, model, "relay-drain-characterization")
 
 	// Attach before processing so the relay opens as soon as
 	// status=running arrives.
@@ -1699,11 +1678,11 @@ func TestSubscribeRelayEstablishedMidStream(t *testing.T) {
 	// call) involves multiple DB round-trips that can be slow under
 	// load.
 	ctx := testutil.Context(t, testutil.WaitSuperLong)
-	user, org, model := seedChatDependencies(ctx, t, db)
+	user, org, model := seedChatDependencies(t, db)
 	setOpenAIProviderBaseURL(ctx, t, db, openAIURL)
 
 	// Create the chat in waiting state.
-	chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "mid-stream-relay")
+	chat := seedWaitingChat(t, db, org.ID, user, model, "mid-stream-relay")
 
 	// Subscribe from the subscriber replica while the chat is idle.
 	_, events, subCancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)

enterprise/coderd/x/chatd/usagelimit_test.go

@@ -1,11 +1,13 @@
 package chatd_test
 
 import (
+	"database/sql"
 	"encoding/json"
 	"testing"
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/sqlc-dev/pqtype"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/v2/coderd/database"
@@ -86,28 +88,14 @@ func TestResolveUsageLimitStatus_OrgScoped(t *testing.T) {
 	require.NoError(t, err)
 
 	// We need a chat provider + model config for inserting chats.
-	_, err = db.InsertChatProvider(ctx, database.InsertChatProviderParams{
-		Provider:             "openai",
-		DisplayName:          "openai",
-		APIKey:               "test-key",
-		CreatedBy:            uuid.NullUUID{UUID: user.ID, Valid: true},
-		Enabled:              true,
-		CentralApiKeyEnabled: true,
+	_ = dbgen.ChatProvider(t, db, database.ChatProvider{
+		CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
 	})
-	require.NoError(t, err)
-	modelConfig, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
-		Provider:             "openai",
-		Model:                "gpt-4o-mini",
-		DisplayName:          "Test Model",
-		CreatedBy:            uuid.NullUUID{UUID: user.ID, Valid: true},
-		UpdatedBy:            uuid.NullUUID{UUID: user.ID, Valid: true},
-		Enabled:              true,
-		IsDefault:            true,
-		ContextLimit:         128000,
-		CompressionThreshold: 70,
-		Options:              json.RawMessage(`{}`),
+	modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
+		CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
+		UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
+		IsDefault: true,
 	})
-	require.NoError(t, err)
 
 	now := time.Now().UTC()
 
@@ -115,38 +103,25 @@ func TestResolveUsageLimitStatus_OrgScoped(t *testing.T) {
 	// given org and inserts a single message with the specified cost.
 	insertChatWithSpend := func(t *testing.T, ownerID, orgID, modelCfgID uuid.UUID, costMicros int64) {
 		t.Helper()
-		tctx := testutil.Context(t, testutil.WaitLong)
-		c, err := db.InsertChat(tctx, database.InsertChatParams{
+		c := dbgen.Chat(t, db, database.Chat{
 			OrganizationID:    orgID,
 			OwnerID:           ownerID,
 			LastModelConfigID: modelCfgID,
 			Title:             "test chat",
-			Status:            database.ChatStatusWaiting,
-			ClientType:        database.ChatClientTypeUi,
-			MCPServerIDs:      []uuid.UUID{},
 		})
-		require.NoError(t, err)
-		_, err = db.InsertChatMessages(tctx, database.InsertChatMessagesParams{
-			ChatID:              c.ID,
-			CreatedBy:           []uuid.UUID{uuid.Nil},
-			ModelConfigID:       []uuid.UUID{modelCfgID},
-			Role:                []database.ChatMessageRole{database.ChatMessageRoleAssistant},
-			Content:             []string{`[{"type":"text","text":"hello"}]`},
-			ContentVersion:      []int16{1},
-			Visibility:          []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
-			InputTokens:         []int64{100},
-			OutputTokens:        []int64{50},
-			TotalTokens:         []int64{150},
-			ReasoningTokens:     []int64{0},
-			CacheCreationTokens: []int64{0},
-			CacheReadTokens:     []int64{0},
-			ContextLimit:        []int64{128000},
-			Compressed:          []bool{false},
-			TotalCostMicros:     []int64{costMicros},
-			RuntimeMs:           []int64{500},
-			ProviderResponseID:  []string{uuid.NewString()},
+		_ = dbgen.ChatMessage(t, db, database.ChatMessage{
+			ChatID:             c.ID,
+			ModelConfigID:      uuid.NullUUID{UUID: modelCfgID, Valid: true},
+			Role:               database.ChatMessageRoleAssistant,
+			Content:            pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"hello"}]`), Valid: true},
+			InputTokens:        sql.NullInt64{Int64: 100, Valid: true},
+			OutputTokens:       sql.NullInt64{Int64: 50, Valid: true},
+			TotalTokens:        sql.NullInt64{Int64: 150, Valid: true},
+			ContextLimit:       sql.NullInt64{Int64: 128000, Valid: true},
+			TotalCostMicros:    sql.NullInt64{Int64: costMicros, Valid: true},
+			RuntimeMs:          sql.NullInt64{Int64: 500, Valid: true},
+			ProviderResponseID: sql.NullString{String: uuid.NewString(), Valid: true},
 		})
-		require.NoError(t, err)
 	}
 
 	t.Run("OrgA_gets_orgA_limit", func(t *testing.T) {