mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat(coderd): refactors github pr sync functionality (#22715)
- Adds `_API_BASE_URL` to `CODER_EXTERNAL_AUTH_CONFIG_` - Extracts and refactors existing GitHub PR sync logic to new packages `coderd/gitsync` and `coderd/externalauth/gitprovider` - Associated wiring and tests Created using Opus 4.6
This commit is contained in:
@@ -0,0 +1,230 @@
|
||||
package gitsync
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/externalauth/gitprovider"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
const (
|
||||
// DiffStatusTTL is how long a successfully refreshed
|
||||
// diff status remains fresh before becoming stale again.
|
||||
DiffStatusTTL = 120 * time.Second
|
||||
)
|
||||
|
||||
// ProviderResolver maps a git remote origin to the gitprovider
|
||||
// that handles it. Returns nil if no provider matches.
|
||||
type ProviderResolver func(origin string) gitprovider.Provider
|
||||
|
||||
var ErrNoTokenAvailable error = errors.New("no token available")
|
||||
|
||||
// TokenResolver obtains the user's git access token for a given
|
||||
// remote origin. Should return nil if no token is available, in
|
||||
// which case ErrNoTokenAvailable will be returned.
|
||||
type TokenResolver func(
|
||||
ctx context.Context,
|
||||
userID uuid.UUID,
|
||||
origin string,
|
||||
) (*string, error)
|
||||
|
||||
// Refresher contains the stateless business logic for fetching
|
||||
// fresh PR data from a git provider given a stale
|
||||
// database.ChatDiffStatus row.
|
||||
type Refresher struct {
|
||||
providers ProviderResolver
|
||||
tokens TokenResolver
|
||||
logger slog.Logger
|
||||
clock quartz.Clock
|
||||
}
|
||||
|
||||
// NewRefresher creates a Refresher with the given dependency
|
||||
// functions.
|
||||
func NewRefresher(
|
||||
providers ProviderResolver,
|
||||
tokens TokenResolver,
|
||||
logger slog.Logger,
|
||||
clock quartz.Clock,
|
||||
) *Refresher {
|
||||
return &Refresher{
|
||||
providers: providers,
|
||||
tokens: tokens,
|
||||
logger: logger,
|
||||
clock: clock,
|
||||
}
|
||||
}
|
||||
|
||||
// RefreshRequest pairs a stale row with the chat owner who
|
||||
// holds the git token needed for API calls.
|
||||
type RefreshRequest struct {
|
||||
Row database.ChatDiffStatus
|
||||
OwnerID uuid.UUID
|
||||
}
|
||||
|
||||
// RefreshResult is the outcome for a single row.
|
||||
// - Params != nil, Error == nil → success, caller should upsert.
|
||||
// - Params == nil, Error == nil → no PR yet, caller should skip.
|
||||
// - Params == nil, Error != nil → row-level failure.
|
||||
type RefreshResult struct {
|
||||
Request RefreshRequest
|
||||
Params *database.UpsertChatDiffStatusParams
|
||||
Error error
|
||||
}
|
||||
|
||||
// groupKey identifies a unique (owner, origin) pair so that
|
||||
// provider and token resolution happen once per group.
|
||||
type groupKey struct {
|
||||
ownerID uuid.UUID
|
||||
origin string
|
||||
}
|
||||
|
||||
// Refresh fetches fresh PR data for a batch of stale rows.
|
||||
// Rows are grouped internally by (ownerID, origin) so that
|
||||
// provider and token resolution happen once per group. A
|
||||
// top-level error is returned only when the entire batch
|
||||
// fails catastrophically. Per-row outcomes are in the
|
||||
// returned RefreshResult slice (one per input request, same
|
||||
// order).
|
||||
func (r *Refresher) Refresh(
|
||||
ctx context.Context,
|
||||
requests []RefreshRequest,
|
||||
) ([]RefreshResult, error) {
|
||||
results := make([]RefreshResult, len(requests))
|
||||
for i, req := range requests {
|
||||
results[i].Request = req
|
||||
}
|
||||
|
||||
// Group request indices by (ownerID, origin).
|
||||
groups := make(map[groupKey][]int)
|
||||
for i, req := range requests {
|
||||
key := groupKey{
|
||||
ownerID: req.OwnerID,
|
||||
origin: req.Row.GitRemoteOrigin,
|
||||
}
|
||||
groups[key] = append(groups[key], i)
|
||||
}
|
||||
|
||||
for key, indices := range groups {
|
||||
provider := r.providers(key.origin)
|
||||
if provider == nil {
|
||||
err := xerrors.Errorf("no provider for origin %q", key.origin)
|
||||
for _, i := range indices {
|
||||
results[i].Error = err
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
token, err := r.tokens(ctx, key.ownerID, key.origin)
|
||||
if err != nil {
|
||||
err = xerrors.Errorf("resolve token: %w", err)
|
||||
} else if token == nil || len(*token) == 0 {
|
||||
err = ErrNoTokenAvailable
|
||||
}
|
||||
if err != nil {
|
||||
for _, i := range indices {
|
||||
results[i].Error = err
|
||||
}
|
||||
continue
|
||||
}
|
||||
// This is technically unnecessary but kept here as a future molly-guard.
|
||||
if token == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for i, idx := range indices {
|
||||
req := requests[idx]
|
||||
params, err := r.refreshOne(ctx, provider, *token, req.Row)
|
||||
results[idx] = RefreshResult{Request: req, Params: params, Error: err}
|
||||
|
||||
// If rate-limited, skip remaining rows in this group.
|
||||
var rlErr *gitprovider.RateLimitError
|
||||
if errors.As(err, &rlErr) {
|
||||
for _, remaining := range indices[i+1:] {
|
||||
results[remaining] = RefreshResult{
|
||||
Request: requests[remaining],
|
||||
Error: fmt.Errorf("skipped: %w", rlErr),
|
||||
}
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// refreshOne processes a single row using an already-resolved
|
||||
// provider and token. This is the old Refresh logic, unchanged.
|
||||
func (r *Refresher) refreshOne(
|
||||
ctx context.Context,
|
||||
provider gitprovider.Provider,
|
||||
token string,
|
||||
row database.ChatDiffStatus,
|
||||
) (*database.UpsertChatDiffStatusParams, error) {
|
||||
var ref gitprovider.PRRef
|
||||
var prURL string
|
||||
|
||||
if row.Url.Valid && row.Url.String != "" {
|
||||
// Row already has a PR URL — parse it directly.
|
||||
parsed, ok := provider.ParsePullRequestURL(row.Url.String)
|
||||
if !ok {
|
||||
return nil, xerrors.Errorf("parse pull request URL %q", row.Url.String)
|
||||
}
|
||||
ref = parsed
|
||||
prURL = row.Url.String
|
||||
} else {
|
||||
// No PR URL — resolve owner/repo from the remote origin,
|
||||
// then look up the open PR for this branch.
|
||||
owner, repo, _, ok := provider.ParseRepositoryOrigin(row.GitRemoteOrigin)
|
||||
if !ok {
|
||||
return nil, xerrors.Errorf("parse repository origin %q", row.GitRemoteOrigin)
|
||||
}
|
||||
|
||||
resolved, err := provider.ResolveBranchPullRequest(ctx, token, gitprovider.BranchRef{
|
||||
Owner: owner,
|
||||
Repo: repo,
|
||||
Branch: row.GitBranch,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("resolve branch pull request: %w", err)
|
||||
}
|
||||
if resolved == nil {
|
||||
// No PR exists yet for this branch.
|
||||
return nil, nil
|
||||
}
|
||||
ref = *resolved
|
||||
prURL = provider.BuildPullRequestURL(ref)
|
||||
}
|
||||
|
||||
status, err := provider.FetchPullRequestStatus(ctx, token, ref)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("fetch pull request status: %w", err)
|
||||
}
|
||||
|
||||
now := r.clock.Now().UTC()
|
||||
params := &database.UpsertChatDiffStatusParams{
|
||||
ChatID: row.ChatID,
|
||||
Url: sql.NullString{String: prURL, Valid: prURL != ""},
|
||||
PullRequestState: sql.NullString{
|
||||
String: string(status.State),
|
||||
Valid: status.State != "",
|
||||
},
|
||||
ChangesRequested: status.ChangesRequested,
|
||||
Additions: status.DiffStats.Additions,
|
||||
Deletions: status.DiffStats.Deletions,
|
||||
ChangedFiles: status.DiffStats.ChangedFiles,
|
||||
RefreshedAt: now,
|
||||
StaleAt: now.Add(DiffStatusTTL),
|
||||
}
|
||||
|
||||
return params, nil
|
||||
}
|
||||
@@ -0,0 +1,775 @@
|
||||
package gitsync_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/externalauth/gitprovider"
|
||||
"github.com/coder/coder/v2/coderd/gitsync"
|
||||
"github.com/coder/coder/v2/coderd/util/ptr"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
// mockProvider implements gitprovider.Provider with function fields
|
||||
// so each test can wire only the methods it needs. Any method left
|
||||
// nil panics with "unexpected call".
|
||||
type mockProvider struct {
|
||||
fetchPullRequestStatus func(ctx context.Context, token string, ref gitprovider.PRRef) (*gitprovider.PRStatus, error)
|
||||
resolveBranchPR func(ctx context.Context, token string, ref gitprovider.BranchRef) (*gitprovider.PRRef, error)
|
||||
fetchPullRequestDiff func(ctx context.Context, token string, ref gitprovider.PRRef) (string, error)
|
||||
fetchBranchDiff func(ctx context.Context, token string, ref gitprovider.BranchRef) (string, error)
|
||||
parseRepositoryOrigin func(raw string) (string, string, string, bool)
|
||||
parsePullRequestURL func(raw string) (gitprovider.PRRef, bool)
|
||||
normalizePullRequestURL func(raw string) string
|
||||
buildBranchURL func(owner, repo, branch string) string
|
||||
buildRepositoryURL func(owner, repo string) string
|
||||
buildPullRequestURL func(ref gitprovider.PRRef) string
|
||||
}
|
||||
|
||||
func (m *mockProvider) FetchPullRequestStatus(ctx context.Context, token string, ref gitprovider.PRRef) (*gitprovider.PRStatus, error) {
|
||||
if m.fetchPullRequestStatus == nil {
|
||||
panic("unexpected call to FetchPullRequestStatus")
|
||||
}
|
||||
return m.fetchPullRequestStatus(ctx, token, ref)
|
||||
}
|
||||
|
||||
func (m *mockProvider) ResolveBranchPullRequest(ctx context.Context, token string, ref gitprovider.BranchRef) (*gitprovider.PRRef, error) {
|
||||
if m.resolveBranchPR == nil {
|
||||
panic("unexpected call to ResolveBranchPullRequest")
|
||||
}
|
||||
return m.resolveBranchPR(ctx, token, ref)
|
||||
}
|
||||
|
||||
func (m *mockProvider) FetchPullRequestDiff(ctx context.Context, token string, ref gitprovider.PRRef) (string, error) {
|
||||
if m.fetchPullRequestDiff == nil {
|
||||
panic("unexpected call to FetchPullRequestDiff")
|
||||
}
|
||||
return m.fetchPullRequestDiff(ctx, token, ref)
|
||||
}
|
||||
|
||||
func (m *mockProvider) FetchBranchDiff(ctx context.Context, token string, ref gitprovider.BranchRef) (string, error) {
|
||||
if m.fetchBranchDiff == nil {
|
||||
panic("unexpected call to FetchBranchDiff")
|
||||
}
|
||||
return m.fetchBranchDiff(ctx, token, ref)
|
||||
}
|
||||
|
||||
func (m *mockProvider) ParseRepositoryOrigin(raw string) (string, string, string, bool) {
|
||||
if m.parseRepositoryOrigin == nil {
|
||||
panic("unexpected call to ParseRepositoryOrigin")
|
||||
}
|
||||
return m.parseRepositoryOrigin(raw)
|
||||
}
|
||||
|
||||
func (m *mockProvider) ParsePullRequestURL(raw string) (gitprovider.PRRef, bool) {
|
||||
if m.parsePullRequestURL == nil {
|
||||
panic("unexpected call to ParsePullRequestURL")
|
||||
}
|
||||
return m.parsePullRequestURL(raw)
|
||||
}
|
||||
|
||||
func (m *mockProvider) NormalizePullRequestURL(raw string) string {
|
||||
if m.normalizePullRequestURL == nil {
|
||||
panic("unexpected call to NormalizePullRequestURL")
|
||||
}
|
||||
return m.normalizePullRequestURL(raw)
|
||||
}
|
||||
|
||||
func (m *mockProvider) BuildBranchURL(owner, repo, branch string) string {
|
||||
if m.buildBranchURL == nil {
|
||||
panic("unexpected call to BuildBranchURL")
|
||||
}
|
||||
return m.buildBranchURL(owner, repo, branch)
|
||||
}
|
||||
|
||||
func (m *mockProvider) BuildRepositoryURL(owner, repo string) string {
|
||||
if m.buildRepositoryURL == nil {
|
||||
panic("unexpected call to BuildRepositoryURL")
|
||||
}
|
||||
return m.buildRepositoryURL(owner, repo)
|
||||
}
|
||||
|
||||
func (m *mockProvider) BuildPullRequestURL(ref gitprovider.PRRef) string {
|
||||
if m.buildPullRequestURL == nil {
|
||||
panic("unexpected call to BuildPullRequestURL")
|
||||
}
|
||||
return m.buildPullRequestURL(ref)
|
||||
}
|
||||
|
||||
func TestRefresher_WithPRURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
mp := &mockProvider{
|
||||
parsePullRequestURL: func(raw string) (gitprovider.PRRef, bool) {
|
||||
return gitprovider.PRRef{Owner: "org", Repo: "repo", Number: 42}, true
|
||||
},
|
||||
fetchPullRequestStatus: func(_ context.Context, _ string, _ gitprovider.PRRef) (*gitprovider.PRStatus, error) {
|
||||
return &gitprovider.PRStatus{
|
||||
State: gitprovider.PRStateOpen,
|
||||
DiffStats: gitprovider.DiffStats{
|
||||
Additions: 10,
|
||||
Deletions: 5,
|
||||
ChangedFiles: 3,
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
return ptr.Ref("test-token"), nil
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
chatID := uuid.New()
|
||||
row := database.ChatDiffStatus{
|
||||
ChatID: chatID,
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/42", Valid: true},
|
||||
GitRemoteOrigin: "https://github.com/org/repo",
|
||||
GitBranch: "feature",
|
||||
}
|
||||
|
||||
ownerID := uuid.New()
|
||||
results, err := r.Refresh(context.Background(), []gitsync.RefreshRequest{
|
||||
{Row: row, OwnerID: ownerID},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
res := results[0]
|
||||
|
||||
require.NoError(t, res.Error)
|
||||
require.NotNil(t, res.Params)
|
||||
|
||||
assert.Equal(t, chatID, res.Params.ChatID)
|
||||
assert.Equal(t, "open", res.Params.PullRequestState.String)
|
||||
assert.True(t, res.Params.PullRequestState.Valid)
|
||||
assert.Equal(t, int32(10), res.Params.Additions)
|
||||
assert.Equal(t, int32(5), res.Params.Deletions)
|
||||
assert.Equal(t, int32(3), res.Params.ChangedFiles)
|
||||
|
||||
// StaleAt should be ~120s after RefreshedAt.
|
||||
diff := res.Params.StaleAt.Sub(res.Params.RefreshedAt)
|
||||
assert.InDelta(t, 120, diff.Seconds(), 5)
|
||||
}
|
||||
|
||||
func TestRefresher_BranchResolvesToPR(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
mp := &mockProvider{
|
||||
parseRepositoryOrigin: func(_ string) (string, string, string, bool) {
|
||||
return "org", "repo", "https://github.com/org/repo", true
|
||||
},
|
||||
resolveBranchPR: func(_ context.Context, _ string, _ gitprovider.BranchRef) (*gitprovider.PRRef, error) {
|
||||
return &gitprovider.PRRef{Owner: "org", Repo: "repo", Number: 7}, nil
|
||||
},
|
||||
fetchPullRequestStatus: func(_ context.Context, _ string, _ gitprovider.PRRef) (*gitprovider.PRStatus, error) {
|
||||
return &gitprovider.PRStatus{State: gitprovider.PRStateOpen}, nil
|
||||
},
|
||||
buildPullRequestURL: func(_ gitprovider.PRRef) string {
|
||||
return "https://github.com/org/repo/pull/7"
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
return ptr.Ref("test-token"), nil
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
row := database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{},
|
||||
GitRemoteOrigin: "https://github.com/org/repo",
|
||||
GitBranch: "feature",
|
||||
}
|
||||
|
||||
ownerID := uuid.New()
|
||||
results, err := r.Refresh(context.Background(), []gitsync.RefreshRequest{
|
||||
{Row: row, OwnerID: ownerID},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
res := results[0]
|
||||
|
||||
require.NoError(t, res.Error)
|
||||
require.NotNil(t, res.Params)
|
||||
|
||||
assert.Contains(t, res.Params.Url.String, "pull/7")
|
||||
assert.True(t, res.Params.Url.Valid)
|
||||
assert.Equal(t, "open", res.Params.PullRequestState.String)
|
||||
}
|
||||
|
||||
func TestRefresher_BranchNoPRYet(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
mp := &mockProvider{
|
||||
parseRepositoryOrigin: func(_ string) (string, string, string, bool) {
|
||||
return "org", "repo", "https://github.com/org/repo", true
|
||||
},
|
||||
resolveBranchPR: func(_ context.Context, _ string, _ gitprovider.BranchRef) (*gitprovider.PRRef, error) {
|
||||
return nil, nil
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
return ptr.Ref("test-token"), nil
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
row := database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{},
|
||||
GitRemoteOrigin: "https://github.com/org/repo",
|
||||
GitBranch: "feature",
|
||||
}
|
||||
|
||||
ownerID := uuid.New()
|
||||
results, err := r.Refresh(context.Background(), []gitsync.RefreshRequest{
|
||||
{Row: row, OwnerID: ownerID},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
res := results[0]
|
||||
|
||||
assert.NoError(t, res.Error)
|
||||
assert.Nil(t, res.Params)
|
||||
}
|
||||
|
||||
func TestRefresher_NoProviderForOrigin(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return nil }
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
return ptr.Ref("test-token"), nil
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
row := database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://example.com/pr/1", Valid: true},
|
||||
GitRemoteOrigin: "https://example.com/org/repo",
|
||||
GitBranch: "feature",
|
||||
}
|
||||
|
||||
ownerID := uuid.New()
|
||||
results, err := r.Refresh(context.Background(), []gitsync.RefreshRequest{
|
||||
{Row: row, OwnerID: ownerID},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
res := results[0]
|
||||
|
||||
assert.Nil(t, res.Params)
|
||||
require.Error(t, res.Error)
|
||||
assert.Contains(t, res.Error.Error(), "no provider")
|
||||
}
|
||||
|
||||
func TestRefresher_TokenResolutionFails(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var fetchCalled atomic.Bool
|
||||
mp := &mockProvider{
|
||||
fetchPullRequestStatus: func(_ context.Context, _ string, _ gitprovider.PRRef) (*gitprovider.PRStatus, error) {
|
||||
fetchCalled.Store(true)
|
||||
return nil, errors.New("should not be called")
|
||||
},
|
||||
parsePullRequestURL: func(_ string) (gitprovider.PRRef, bool) {
|
||||
return gitprovider.PRRef{Owner: "org", Repo: "repo", Number: 1}, true
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
return nil, errors.New("token lookup failed")
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
row := database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/1", Valid: true},
|
||||
GitRemoteOrigin: "https://github.com/org/repo",
|
||||
GitBranch: "feature",
|
||||
}
|
||||
|
||||
ownerID := uuid.New()
|
||||
results, err := r.Refresh(context.Background(), []gitsync.RefreshRequest{
|
||||
{Row: row, OwnerID: ownerID},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
res := results[0]
|
||||
|
||||
assert.Nil(t, res.Params)
|
||||
require.Error(t, res.Error)
|
||||
assert.False(t, fetchCalled.Load(), "FetchPullRequestStatus should not be called when token resolution fails")
|
||||
}
|
||||
|
||||
func TestRefresher_EmptyToken(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
mp := &mockProvider{}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
return ptr.Ref(""), nil
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
row := database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/1", Valid: true},
|
||||
GitRemoteOrigin: "https://github.com/org/repo",
|
||||
GitBranch: "feature",
|
||||
}
|
||||
|
||||
ownerID := uuid.New()
|
||||
results, err := r.Refresh(context.Background(), []gitsync.RefreshRequest{
|
||||
{Row: row, OwnerID: ownerID},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
res := results[0]
|
||||
|
||||
assert.Nil(t, res.Params)
|
||||
require.ErrorIs(t, res.Error, gitsync.ErrNoTokenAvailable)
|
||||
}
|
||||
|
||||
func TestRefresher_ProviderFetchFails(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
mp := &mockProvider{
|
||||
parsePullRequestURL: func(_ string) (gitprovider.PRRef, bool) {
|
||||
return gitprovider.PRRef{Owner: "org", Repo: "repo", Number: 42}, true
|
||||
},
|
||||
fetchPullRequestStatus: func(_ context.Context, _ string, _ gitprovider.PRRef) (*gitprovider.PRStatus, error) {
|
||||
return nil, errors.New("api error")
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
return ptr.Ref("test-token"), nil
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
row := database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/42", Valid: true},
|
||||
GitRemoteOrigin: "https://github.com/org/repo",
|
||||
GitBranch: "feature",
|
||||
}
|
||||
|
||||
ownerID := uuid.New()
|
||||
results, err := r.Refresh(context.Background(), []gitsync.RefreshRequest{
|
||||
{Row: row, OwnerID: ownerID},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
res := results[0]
|
||||
|
||||
assert.Nil(t, res.Params)
|
||||
require.Error(t, res.Error)
|
||||
assert.Contains(t, res.Error.Error(), "api error")
|
||||
}
|
||||
|
||||
func TestRefresher_PRURLParseFailure(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
mp := &mockProvider{
|
||||
parsePullRequestURL: func(_ string) (gitprovider.PRRef, bool) {
|
||||
return gitprovider.PRRef{}, false
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
return ptr.Ref("test-token"), nil
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
row := database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/not-a-pr", Valid: true},
|
||||
GitRemoteOrigin: "https://github.com/org/repo",
|
||||
GitBranch: "feature",
|
||||
}
|
||||
|
||||
ownerID := uuid.New()
|
||||
results, err := r.Refresh(context.Background(), []gitsync.RefreshRequest{
|
||||
{Row: row, OwnerID: ownerID},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
res := results[0]
|
||||
|
||||
assert.Nil(t, res.Params)
|
||||
require.Error(t, res.Error)
|
||||
}
|
||||
|
||||
func TestRefresher_BatchGroupsByOwnerAndOrigin(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
mp := &mockProvider{
|
||||
parsePullRequestURL: func(_ string) (gitprovider.PRRef, bool) {
|
||||
return gitprovider.PRRef{Owner: "org", Repo: "repo", Number: 1}, true
|
||||
},
|
||||
fetchPullRequestStatus: func(_ context.Context, _ string, _ gitprovider.PRRef) (*gitprovider.PRStatus, error) {
|
||||
return &gitprovider.PRStatus{State: gitprovider.PRStateOpen}, nil
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
|
||||
var tokenCalls atomic.Int32
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
tokenCalls.Add(1)
|
||||
return ptr.Ref("test-token"), nil
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
ownerID := uuid.New()
|
||||
originA := "https://github.com/org/repo"
|
||||
originB := "https://gitlab.com/org/repo"
|
||||
|
||||
requests := []gitsync.RefreshRequest{
|
||||
{
|
||||
Row: database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/1", Valid: true},
|
||||
GitRemoteOrigin: originA,
|
||||
GitBranch: "feature-1",
|
||||
},
|
||||
OwnerID: ownerID,
|
||||
},
|
||||
{
|
||||
Row: database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/1", Valid: true},
|
||||
GitRemoteOrigin: originA,
|
||||
GitBranch: "feature-2",
|
||||
},
|
||||
OwnerID: ownerID,
|
||||
},
|
||||
{
|
||||
Row: database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://gitlab.com/org/repo/pull/1", Valid: true},
|
||||
GitRemoteOrigin: originB,
|
||||
GitBranch: "feature-3",
|
||||
},
|
||||
OwnerID: ownerID,
|
||||
},
|
||||
}
|
||||
|
||||
results, err := r.Refresh(context.Background(), requests)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 3)
|
||||
|
||||
for i, res := range results {
|
||||
require.NoError(t, res.Error, "result[%d] should not have an error", i)
|
||||
require.NotNil(t, res.Params, "result[%d] should have params", i)
|
||||
}
|
||||
|
||||
// Two distinct (ownerID, origin) groups → exactly 2 token
|
||||
// resolution calls.
|
||||
assert.Equal(t, int32(2), tokenCalls.Load(),
|
||||
"TokenResolver should be called once per (owner, origin) group")
|
||||
}
|
||||
|
||||
func TestRefresher_UsesInjectedClock(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
fixedTime := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC)
|
||||
mClock.Set(fixedTime)
|
||||
|
||||
mp := &mockProvider{
|
||||
parsePullRequestURL: func(raw string) (gitprovider.PRRef, bool) {
|
||||
return gitprovider.PRRef{Owner: "org", Repo: "repo", Number: 42}, true
|
||||
},
|
||||
fetchPullRequestStatus: func(_ context.Context, _ string, _ gitprovider.PRRef) (*gitprovider.PRStatus, error) {
|
||||
return &gitprovider.PRStatus{
|
||||
State: gitprovider.PRStateOpen,
|
||||
DiffStats: gitprovider.DiffStats{
|
||||
Additions: 10,
|
||||
Deletions: 5,
|
||||
ChangedFiles: 3,
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
return ptr.Ref("test-token"), nil
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), mClock)
|
||||
|
||||
chatID := uuid.New()
|
||||
row := database.ChatDiffStatus{
|
||||
ChatID: chatID,
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/42", Valid: true},
|
||||
GitRemoteOrigin: "https://github.com/org/repo",
|
||||
GitBranch: "feature",
|
||||
}
|
||||
|
||||
ownerID := uuid.New()
|
||||
results, err := r.Refresh(context.Background(), []gitsync.RefreshRequest{
|
||||
{Row: row, OwnerID: ownerID},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
res := results[0]
|
||||
|
||||
require.NoError(t, res.Error)
|
||||
require.NotNil(t, res.Params)
|
||||
|
||||
// The mock clock is deterministic, so times must be exact.
|
||||
assert.Equal(t, fixedTime, res.Params.RefreshedAt)
|
||||
assert.Equal(t, fixedTime.Add(gitsync.DiffStatusTTL), res.Params.StaleAt)
|
||||
}
|
||||
|
||||
func TestRefresher_RateLimitSkipsRemainingInGroup(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var callCount atomic.Int32
|
||||
|
||||
mp := &mockProvider{
|
||||
parsePullRequestURL: func(raw string) (gitprovider.PRRef, bool) {
|
||||
var num int
|
||||
switch {
|
||||
case strings.HasSuffix(raw, "/pull/1"):
|
||||
num = 1
|
||||
case strings.HasSuffix(raw, "/pull/2"):
|
||||
num = 2
|
||||
case strings.HasSuffix(raw, "/pull/3"):
|
||||
num = 3
|
||||
default:
|
||||
return gitprovider.PRRef{}, false
|
||||
}
|
||||
return gitprovider.PRRef{Owner: "org", Repo: "repo", Number: num}, true
|
||||
},
|
||||
fetchPullRequestStatus: func(_ context.Context, _ string, ref gitprovider.PRRef) (*gitprovider.PRStatus, error) {
|
||||
call := callCount.Add(1)
|
||||
switch call {
|
||||
case 1:
|
||||
// First call succeeds.
|
||||
return &gitprovider.PRStatus{
|
||||
State: gitprovider.PRStateOpen,
|
||||
DiffStats: gitprovider.DiffStats{
|
||||
Additions: 5,
|
||||
Deletions: 2,
|
||||
ChangedFiles: 1,
|
||||
},
|
||||
}, nil
|
||||
case 2:
|
||||
// Second call hits rate limit.
|
||||
return nil, &gitprovider.RateLimitError{
|
||||
RetryAfter: time.Now().Add(60 * time.Second),
|
||||
}
|
||||
default:
|
||||
// Third call should never happen.
|
||||
t.Fatal("FetchPullRequestStatus called more than 2 times")
|
||||
return nil, nil
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
tokens := func(_ context.Context, _ uuid.UUID, _ string) (*string, error) {
|
||||
return ptr.Ref("test-token"), nil
|
||||
}
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
ownerID := uuid.New()
|
||||
origin := "https://github.com/org/repo"
|
||||
|
||||
requests := []gitsync.RefreshRequest{
|
||||
{
|
||||
Row: database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/1", Valid: true},
|
||||
GitRemoteOrigin: origin,
|
||||
GitBranch: "feat-1",
|
||||
},
|
||||
OwnerID: ownerID,
|
||||
},
|
||||
{
|
||||
Row: database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/2", Valid: true},
|
||||
GitRemoteOrigin: origin,
|
||||
GitBranch: "feat-2",
|
||||
},
|
||||
OwnerID: ownerID,
|
||||
},
|
||||
{
|
||||
Row: database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/3", Valid: true},
|
||||
GitRemoteOrigin: origin,
|
||||
GitBranch: "feat-3",
|
||||
},
|
||||
OwnerID: ownerID,
|
||||
},
|
||||
}
|
||||
|
||||
results, err := r.Refresh(context.Background(), requests)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 3)
|
||||
|
||||
// Row 0: success.
|
||||
assert.NoError(t, results[0].Error)
|
||||
assert.NotNil(t, results[0].Params)
|
||||
|
||||
// Row 1: rate-limited.
|
||||
require.Error(t, results[1].Error)
|
||||
var rlErr1 *gitprovider.RateLimitError
|
||||
assert.True(t, errors.As(results[1].Error, &rlErr1),
|
||||
"result[1] error should be *RateLimitError")
|
||||
|
||||
// Row 2: skipped due to rate limit.
|
||||
require.Error(t, results[2].Error)
|
||||
var rlErr2 *gitprovider.RateLimitError
|
||||
assert.True(t, errors.As(results[2].Error, &rlErr2),
|
||||
"result[2] error should wrap *RateLimitError")
|
||||
assert.Contains(t, results[2].Error.Error(), "skipped")
|
||||
|
||||
// Provider should have been called exactly twice.
|
||||
assert.Equal(t, int32(2), callCount.Load(),
|
||||
"FetchPullRequestStatus should be called exactly 2 times")
|
||||
}
|
||||
|
||||
func TestRefresher_CorrectTokenPerOrigin(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var tokenCalls atomic.Int32
|
||||
tokens := func(_ context.Context, _ uuid.UUID, origin string) (*string, error) {
|
||||
tokenCalls.Add(1)
|
||||
switch {
|
||||
case strings.Contains(origin, "github.com"):
|
||||
return ptr.Ref("gh-public-token"), nil
|
||||
case strings.Contains(origin, "ghes.corp.com"):
|
||||
return ptr.Ref("ghe-private-token"), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unexpected origin: %s", origin)
|
||||
}
|
||||
}
|
||||
|
||||
// Track which token each FetchPullRequestStatus call received,
|
||||
// keyed by chat ID. We pass the chat ID through the PRRef.Number
|
||||
// field (unique per request) so FetchPullRequestStatus can
|
||||
// identify which row it's processing.
|
||||
var mu sync.Mutex
|
||||
tokensByPR := make(map[int]string)
|
||||
|
||||
mp := &mockProvider{
|
||||
parsePullRequestURL: func(raw string) (gitprovider.PRRef, bool) {
|
||||
// Extract a unique PR number from the URL to identify
|
||||
// each row inside FetchPullRequestStatus.
|
||||
var num int
|
||||
switch {
|
||||
case strings.HasSuffix(raw, "/pull/1"):
|
||||
num = 1
|
||||
case strings.HasSuffix(raw, "/pull/2"):
|
||||
num = 2
|
||||
case strings.HasSuffix(raw, "/pull/10"):
|
||||
num = 10
|
||||
default:
|
||||
return gitprovider.PRRef{}, false
|
||||
}
|
||||
return gitprovider.PRRef{Owner: "org", Repo: "repo", Number: num}, true
|
||||
},
|
||||
fetchPullRequestStatus: func(_ context.Context, token string, ref gitprovider.PRRef) (*gitprovider.PRStatus, error) {
|
||||
mu.Lock()
|
||||
tokensByPR[ref.Number] = token
|
||||
mu.Unlock()
|
||||
return &gitprovider.PRStatus{State: gitprovider.PRStateOpen}, nil
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(_ string) gitprovider.Provider { return mp }
|
||||
|
||||
r := gitsync.NewRefresher(providers, tokens, slogtest.Make(t, nil), quartz.NewReal())
|
||||
|
||||
ownerID := uuid.New()
|
||||
|
||||
requests := []gitsync.RefreshRequest{
|
||||
{
|
||||
Row: database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/1", Valid: true},
|
||||
GitRemoteOrigin: "https://github.com/org/repo",
|
||||
GitBranch: "feature-1",
|
||||
},
|
||||
OwnerID: ownerID,
|
||||
},
|
||||
{
|
||||
Row: database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://github.com/org/repo/pull/2", Valid: true},
|
||||
GitRemoteOrigin: "https://github.com/org/repo",
|
||||
GitBranch: "feature-2",
|
||||
},
|
||||
OwnerID: ownerID,
|
||||
},
|
||||
{
|
||||
Row: database.ChatDiffStatus{
|
||||
ChatID: uuid.New(),
|
||||
Url: sql.NullString{String: "https://ghes.corp.com/org/repo/pull/10", Valid: true},
|
||||
GitRemoteOrigin: "https://ghes.corp.com/org/repo",
|
||||
GitBranch: "feature-3",
|
||||
},
|
||||
OwnerID: ownerID,
|
||||
},
|
||||
}
|
||||
|
||||
results, err := r.Refresh(context.Background(), requests)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 3)
|
||||
|
||||
for i, res := range results {
|
||||
require.NoError(t, res.Error, "result[%d] should not have an error", i)
|
||||
require.NotNil(t, res.Params, "result[%d] should have params", i)
|
||||
}
|
||||
|
||||
// github.com rows (PR #1 and #2) should use the public token.
|
||||
assert.Equal(t, "gh-public-token", tokensByPR[1],
|
||||
"github.com PR #1 should use gh-public-token")
|
||||
assert.Equal(t, "gh-public-token", tokensByPR[2],
|
||||
"github.com PR #2 should use gh-public-token")
|
||||
|
||||
// ghes.corp.com row (PR #10) should use the GHE token.
|
||||
assert.Equal(t, "ghe-private-token", tokensByPR[10],
|
||||
"ghes.corp.com PR #10 should use ghe-private-token")
|
||||
|
||||
// Token resolution should be called exactly twice — once per
|
||||
// (owner, origin) group.
|
||||
assert.Equal(t, int32(2), tokenCalls.Load(),
|
||||
"TokenResolver should be called once per (owner, origin) group")
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
package gitsync
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
const (
|
||||
// defaultBatchSize is the maximum number of stale rows fetched
|
||||
// per tick.
|
||||
defaultBatchSize int32 = 50
|
||||
|
||||
// defaultInterval is the polling interval between ticks.
|
||||
defaultInterval = 10 * time.Second
|
||||
)
|
||||
|
||||
// Store is the narrow DB interface the Worker needs.
|
||||
type Store interface {
|
||||
AcquireStaleChatDiffStatuses(
|
||||
ctx context.Context, limitVal int32,
|
||||
) ([]database.AcquireStaleChatDiffStatusesRow, error)
|
||||
BackoffChatDiffStatus(
|
||||
ctx context.Context, arg database.BackoffChatDiffStatusParams,
|
||||
) error
|
||||
UpsertChatDiffStatus(
|
||||
ctx context.Context, arg database.UpsertChatDiffStatusParams,
|
||||
) (database.ChatDiffStatus, error)
|
||||
UpsertChatDiffStatusReference(
|
||||
ctx context.Context, arg database.UpsertChatDiffStatusReferenceParams,
|
||||
) (database.ChatDiffStatus, error)
|
||||
GetChatsByOwnerID(
|
||||
ctx context.Context, arg database.GetChatsByOwnerIDParams,
|
||||
) ([]database.Chat, error)
|
||||
}
|
||||
|
||||
// EventPublisher notifies the frontend of diff status changes.
|
||||
type PublishDiffStatusChangeFunc func(ctx context.Context, chatID uuid.UUID) error
|
||||
|
||||
// Worker is a background loop that periodically refreshes stale
|
||||
// chat diff statuses by delegating to a Refresher.
|
||||
type Worker struct {
|
||||
store Store
|
||||
refresher *Refresher
|
||||
publishDiffStatusChangeFn PublishDiffStatusChangeFunc
|
||||
clock quartz.Clock
|
||||
logger slog.Logger
|
||||
batchSize int32
|
||||
interval time.Duration
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
// NewWorker creates a Worker with default batch size and interval.
|
||||
func NewWorker(
|
||||
store Store,
|
||||
refresher *Refresher,
|
||||
publisher PublishDiffStatusChangeFunc,
|
||||
clock quartz.Clock,
|
||||
logger slog.Logger,
|
||||
) *Worker {
|
||||
return &Worker{
|
||||
store: store,
|
||||
refresher: refresher,
|
||||
publishDiffStatusChangeFn: publisher,
|
||||
clock: clock,
|
||||
logger: logger,
|
||||
batchSize: defaultBatchSize,
|
||||
interval: defaultInterval,
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start launches the background loop. It blocks until ctx is
|
||||
// cancelled, then closes w.done.
|
||||
func (w *Worker) Start(ctx context.Context) {
|
||||
defer close(w.done)
|
||||
|
||||
ticker := w.clock.NewTicker(w.interval, "gitsync", "worker")
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
w.tick(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Done returns a channel that is closed when the worker exits.
|
||||
func (w *Worker) Done() <-chan struct{} {
|
||||
return w.done
|
||||
}
|
||||
|
||||
func chatDiffStatusFromRow(row database.AcquireStaleChatDiffStatusesRow) database.ChatDiffStatus {
|
||||
return database.ChatDiffStatus{
|
||||
ChatID: row.ChatID,
|
||||
Url: row.Url,
|
||||
PullRequestState: row.PullRequestState,
|
||||
ChangesRequested: row.ChangesRequested,
|
||||
Additions: row.Additions,
|
||||
Deletions: row.Deletions,
|
||||
ChangedFiles: row.ChangedFiles,
|
||||
RefreshedAt: row.RefreshedAt,
|
||||
StaleAt: row.StaleAt,
|
||||
CreatedAt: row.CreatedAt,
|
||||
UpdatedAt: row.UpdatedAt,
|
||||
GitBranch: row.GitBranch,
|
||||
GitRemoteOrigin: row.GitRemoteOrigin,
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Worker) tick(ctx context.Context) {
|
||||
// Set a context equal to w.interval so that we do not hold up processing due to
|
||||
// random unicorn-related events.
|
||||
ctx, cancel := context.WithTimeout(ctx, w.interval)
|
||||
defer cancel()
|
||||
|
||||
acquiredRows, err := w.store.AcquireStaleChatDiffStatuses(ctx, w.batchSize)
|
||||
if err != nil {
|
||||
w.logger.Warn(ctx, "acquire stale chat diff statuses",
|
||||
slog.Error(err))
|
||||
return
|
||||
}
|
||||
if len(acquiredRows) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Build refresh requests directly from acquired rows.
|
||||
requests := make([]RefreshRequest, 0, len(acquiredRows))
|
||||
for _, row := range acquiredRows {
|
||||
requests = append(requests, RefreshRequest{
|
||||
Row: chatDiffStatusFromRow(row),
|
||||
OwnerID: row.OwnerID,
|
||||
})
|
||||
}
|
||||
|
||||
results, err := w.refresher.Refresh(ctx, requests)
|
||||
if err != nil {
|
||||
w.logger.Warn(ctx, "batch refresh chat diff statuses",
|
||||
slog.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
for _, res := range results {
|
||||
if res.Error != nil {
|
||||
w.logger.Debug(ctx, "refresh chat diff status",
|
||||
slog.F("chat_id", res.Request.Row.ChatID),
|
||||
slog.Error(res.Error))
|
||||
// Back off so the row isn't retried immediately.
|
||||
if err := w.store.BackoffChatDiffStatus(ctx,
|
||||
database.BackoffChatDiffStatusParams{
|
||||
ChatID: res.Request.Row.ChatID,
|
||||
StaleAt: w.clock.Now().UTC().Add(DiffStatusTTL),
|
||||
},
|
||||
); err != nil {
|
||||
w.logger.Warn(ctx, "backoff failed chat diff status",
|
||||
slog.F("chat_id", res.Request.Row.ChatID),
|
||||
slog.Error(err))
|
||||
}
|
||||
continue
|
||||
}
|
||||
if res.Params == nil {
|
||||
// No PR yet — skip.
|
||||
continue
|
||||
}
|
||||
if _, err := w.store.UpsertChatDiffStatus(ctx, *res.Params); err != nil {
|
||||
w.logger.Warn(ctx, "upsert refreshed chat diff status",
|
||||
slog.F("chat_id", res.Request.Row.ChatID),
|
||||
slog.Error(err))
|
||||
continue
|
||||
}
|
||||
if w.publishDiffStatusChangeFn != nil {
|
||||
if err := w.publishDiffStatusChangeFn(ctx, res.Request.Row.ChatID); err != nil {
|
||||
w.logger.Debug(ctx, "publish diff status change",
|
||||
slog.F("chat_id", res.Request.Row.ChatID),
|
||||
slog.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MarkStale persists the git ref on all chats for a workspace,
|
||||
// setting stale_at to the past so the next tick picks them up.
|
||||
// Publishes a diff status event for each affected chat.
|
||||
// Called from workspaceagents handlers. No goroutines spawned.
|
||||
func (w *Worker) MarkStale(
|
||||
ctx context.Context,
|
||||
workspaceID, ownerID uuid.UUID,
|
||||
branch, origin string,
|
||||
) {
|
||||
if branch == "" || origin == "" {
|
||||
return
|
||||
}
|
||||
|
||||
chats, err := w.store.GetChatsByOwnerID(ctx, database.GetChatsByOwnerIDParams{
|
||||
OwnerID: ownerID,
|
||||
})
|
||||
if err != nil {
|
||||
w.logger.Warn(ctx, "list chats for git ref storage",
|
||||
slog.F("workspace_id", workspaceID),
|
||||
slog.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
for _, chat := range filterChatsByWorkspaceID(chats, workspaceID) {
|
||||
_, err := w.store.UpsertChatDiffStatusReference(ctx,
|
||||
database.UpsertChatDiffStatusReferenceParams{
|
||||
ChatID: chat.ID,
|
||||
GitBranch: branch,
|
||||
GitRemoteOrigin: origin,
|
||||
StaleAt: w.clock.Now().Add(-time.Second),
|
||||
Url: sql.NullString{},
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
w.logger.Warn(ctx, "store git ref on chat diff status",
|
||||
slog.F("chat_id", chat.ID),
|
||||
slog.F("workspace_id", workspaceID),
|
||||
slog.Error(err))
|
||||
continue
|
||||
}
|
||||
// Notify the frontend immediately so the UI shows the
|
||||
// branch info even before the worker refreshes PR data.
|
||||
if w.publishDiffStatusChangeFn != nil {
|
||||
if pubErr := w.publishDiffStatusChangeFn(ctx, chat.ID); pubErr != nil {
|
||||
w.logger.Debug(ctx, "publish diff status after mark stale",
|
||||
slog.F("chat_id", chat.ID), slog.Error(pubErr))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// filterChatsByWorkspaceID returns only chats associated with
|
||||
// the given workspace.
|
||||
func filterChatsByWorkspaceID(
|
||||
chats []database.Chat,
|
||||
workspaceID uuid.UUID,
|
||||
) []database.Chat {
|
||||
filtered := make([]database.Chat, 0, len(chats))
|
||||
for _, chat := range chats {
|
||||
if !chat.WorkspaceID.Valid || chat.WorkspaceID.UUID != workspaceID {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, chat)
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
@@ -0,0 +1,744 @@
|
||||
package gitsync_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbgen"
|
||||
"github.com/coder/coder/v2/coderd/database/dbmock"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/coderd/externalauth/gitprovider"
|
||||
"github.com/coder/coder/v2/coderd/gitsync"
|
||||
"github.com/coder/coder/v2/coderd/util/ptr"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
// testRefresherCfg configures newTestRefresher.
|
||||
type testRefresherCfg struct {
|
||||
resolveBranchPR func(context.Context, string, gitprovider.BranchRef) (*gitprovider.PRRef, error)
|
||||
fetchPRStatus func(context.Context, string, gitprovider.PRRef) (*gitprovider.PRStatus, error)
|
||||
}
|
||||
|
||||
type testRefresherOpt func(*testRefresherCfg)
|
||||
|
||||
func withResolveBranchPR(f func(context.Context, string, gitprovider.BranchRef) (*gitprovider.PRRef, error)) testRefresherOpt {
|
||||
return func(c *testRefresherCfg) { c.resolveBranchPR = f }
|
||||
}
|
||||
|
||||
// newTestRefresher creates a Refresher backed by mock
|
||||
// provider/token resolvers. The provider recognises any origin,
|
||||
// resolves branches to a canned PR, and returns a canned PRStatus.
|
||||
func newTestRefresher(t *testing.T, clk quartz.Clock, opts ...testRefresherOpt) *gitsync.Refresher {
|
||||
t.Helper()
|
||||
|
||||
cfg := testRefresherCfg{
|
||||
resolveBranchPR: func(context.Context, string, gitprovider.BranchRef) (*gitprovider.PRRef, error) {
|
||||
return &gitprovider.PRRef{Owner: "o", Repo: "r", Number: 1}, nil
|
||||
},
|
||||
fetchPRStatus: func(context.Context, string, gitprovider.PRRef) (*gitprovider.PRStatus, error) {
|
||||
return &gitprovider.PRStatus{
|
||||
State: gitprovider.PRStateOpen,
|
||||
DiffStats: gitprovider.DiffStats{
|
||||
Additions: 10,
|
||||
Deletions: 3,
|
||||
ChangedFiles: 2,
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(&cfg)
|
||||
}
|
||||
|
||||
prov := &mockProvider{
|
||||
parseRepositoryOrigin: func(string) (string, string, string, bool) {
|
||||
return "owner", "repo", "https://github.com/owner/repo", true
|
||||
},
|
||||
parsePullRequestURL: func(raw string) (gitprovider.PRRef, bool) {
|
||||
return gitprovider.PRRef{Owner: "owner", Repo: "repo", Number: 1}, raw != ""
|
||||
},
|
||||
resolveBranchPR: cfg.resolveBranchPR,
|
||||
fetchPullRequestStatus: cfg.fetchPRStatus,
|
||||
buildPullRequestURL: func(ref gitprovider.PRRef) string {
|
||||
return fmt.Sprintf("https://github.com/%s/%s/pull/%d", ref.Owner, ref.Repo, ref.Number)
|
||||
},
|
||||
}
|
||||
|
||||
providers := func(string) gitprovider.Provider { return prov }
|
||||
tokens := func(context.Context, uuid.UUID, string) (*string, error) {
|
||||
return ptr.Ref("tok"), nil
|
||||
}
|
||||
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
return gitsync.NewRefresher(providers, tokens, logger, clk)
|
||||
}
|
||||
|
||||
// makeAcquiredRow returns an AcquireStaleChatDiffStatusesRow with
|
||||
// a non-empty branch/origin so the Refresher goes through the
|
||||
// branch-resolution path.
|
||||
func makeAcquiredRow(chatID, ownerID uuid.UUID) database.AcquireStaleChatDiffStatusesRow {
|
||||
return database.AcquireStaleChatDiffStatusesRow{
|
||||
ChatID: chatID,
|
||||
GitBranch: "feature",
|
||||
GitRemoteOrigin: "https://github.com/owner/repo",
|
||||
StaleAt: time.Now().Add(-time.Minute),
|
||||
OwnerID: ownerID,
|
||||
}
|
||||
}
|
||||
|
||||
// tickOnce traps the worker's NewTicker call, starts the worker,
|
||||
// fires one tick, waits for it to finish by observing the given
|
||||
// tickDone channel, then shuts the worker down. The tickDone
|
||||
// channel must be closed when the last expected operation in the
|
||||
// tick completes. For tests where the tick does nothing (e.g. 0
|
||||
// stale rows or store error), tickDone should be closed inside
|
||||
// acquireStaleChatDiffStatuses.
|
||||
func tickOnce(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
mClock *quartz.Mock,
|
||||
worker *gitsync.Worker,
|
||||
tickDone <-chan struct{},
|
||||
) {
|
||||
t.Helper()
|
||||
|
||||
trap := mClock.Trap().NewTicker("gitsync", "worker")
|
||||
defer trap.Close()
|
||||
|
||||
workerCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
go worker.Start(workerCtx)
|
||||
|
||||
// Wait for the worker to create its ticker.
|
||||
trap.MustWait(ctx).MustRelease(ctx)
|
||||
|
||||
// Fire one tick. The waiter resolves when the channel receive
|
||||
// completes, not when w.tick() returns, so we use tickDone to
|
||||
// know when to proceed.
|
||||
_, w := mClock.AdvanceNext()
|
||||
w.MustWait(ctx)
|
||||
|
||||
// Wait for the tick's business logic to finish.
|
||||
select {
|
||||
case <-tickDone:
|
||||
case <-ctx.Done():
|
||||
t.Fatal("timed out waiting for tick to complete")
|
||||
}
|
||||
|
||||
cancel()
|
||||
<-worker.Done()
|
||||
}
|
||||
|
||||
func TestWorker_SkipsFreshRows(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
tickDone := make(chan struct{})
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().AcquireStaleChatDiffStatuses(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(context.Context, int32) ([]database.AcquireStaleChatDiffStatusesRow, error) {
|
||||
// No stale rows — tick returns immediately.
|
||||
close(tickDone)
|
||||
return nil, nil
|
||||
})
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
worker := gitsync.NewWorker(store, refresher, nil, mClock, logger)
|
||||
|
||||
tickOnce(ctx, t, mClock, worker, tickDone)
|
||||
}
|
||||
|
||||
func TestWorker_LimitsToNRows(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
var capturedLimit atomic.Int32
|
||||
var upsertCount atomic.Int32
|
||||
ownerID := uuid.New()
|
||||
const numRows = 5
|
||||
tickDone := make(chan struct{})
|
||||
|
||||
rows := make([]database.AcquireStaleChatDiffStatusesRow, numRows)
|
||||
for i := range rows {
|
||||
rows[i] = makeAcquiredRow(uuid.New(), ownerID)
|
||||
}
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().AcquireStaleChatDiffStatuses(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, limitVal int32) ([]database.AcquireStaleChatDiffStatusesRow, error) {
|
||||
capturedLimit.Store(limitVal)
|
||||
return rows, nil
|
||||
})
|
||||
store.EXPECT().UpsertChatDiffStatus(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, arg database.UpsertChatDiffStatusParams) (database.ChatDiffStatus, error) {
|
||||
upsertCount.Add(1)
|
||||
return database.ChatDiffStatus{ChatID: arg.ChatID}, nil
|
||||
}).Times(numRows)
|
||||
|
||||
pub := func(_ context.Context, _ uuid.UUID) error {
|
||||
if upsertCount.Load() == numRows {
|
||||
close(tickDone)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
worker := gitsync.NewWorker(store, refresher, pub, mClock, logger)
|
||||
|
||||
tickOnce(ctx, t, mClock, worker, tickDone)
|
||||
|
||||
// The default batch size is 50.
|
||||
assert.Equal(t, int32(50), capturedLimit.Load())
|
||||
assert.Equal(t, int32(numRows), upsertCount.Load())
|
||||
}
|
||||
|
||||
func TestWorker_RefresherReturnsNilNil_SkipsUpsert(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
chatID := uuid.New()
|
||||
ownerID := uuid.New()
|
||||
|
||||
// When the Refresher returns (nil, nil) the worker skips the
|
||||
// upsert and publish. We signal tickDone from the refresher
|
||||
// mock since that is the last operation before the tick
|
||||
// returns.
|
||||
tickDone := make(chan struct{})
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().AcquireStaleChatDiffStatuses(gomock.Any(), gomock.Any()).
|
||||
Return([]database.AcquireStaleChatDiffStatusesRow{makeAcquiredRow(chatID, ownerID)}, nil)
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
|
||||
// ResolveBranchPullRequest returns nil → Refresher returns
|
||||
// (nil, nil).
|
||||
refresher := newTestRefresher(t, mClock, withResolveBranchPR(
|
||||
func(context.Context, string, gitprovider.BranchRef) (*gitprovider.PRRef, error) {
|
||||
close(tickDone)
|
||||
return nil, nil
|
||||
},
|
||||
))
|
||||
|
||||
worker := gitsync.NewWorker(store, refresher, nil, mClock, logger)
|
||||
|
||||
tickOnce(ctx, t, mClock, worker, tickDone)
|
||||
}
|
||||
|
||||
func TestWorker_RefresherError_BacksOffRow(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
chat1 := uuid.New()
|
||||
chat2 := uuid.New()
|
||||
ownerID := uuid.New()
|
||||
|
||||
var upsertCount atomic.Int32
|
||||
var publishCount atomic.Int32
|
||||
var backoffCount atomic.Int32
|
||||
var mu sync.Mutex
|
||||
var backoffArgs []database.BackoffChatDiffStatusParams
|
||||
tickDone := make(chan struct{})
|
||||
var closeOnce sync.Once
|
||||
|
||||
// Two rows processed: one fails (backoff), one succeeds
|
||||
// (upsert+publish). Both must finish before we close tickDone.
|
||||
var terminalOps atomic.Int32
|
||||
signalIfDone := func() {
|
||||
if terminalOps.Add(1) == 2 {
|
||||
closeOnce.Do(func() { close(tickDone) })
|
||||
}
|
||||
}
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().AcquireStaleChatDiffStatuses(gomock.Any(), gomock.Any()).
|
||||
Return([]database.AcquireStaleChatDiffStatusesRow{
|
||||
makeAcquiredRow(chat1, ownerID),
|
||||
makeAcquiredRow(chat2, ownerID),
|
||||
}, nil)
|
||||
store.EXPECT().BackoffChatDiffStatus(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, arg database.BackoffChatDiffStatusParams) error {
|
||||
backoffCount.Add(1)
|
||||
mu.Lock()
|
||||
backoffArgs = append(backoffArgs, arg)
|
||||
mu.Unlock()
|
||||
signalIfDone()
|
||||
return nil
|
||||
})
|
||||
store.EXPECT().UpsertChatDiffStatus(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, arg database.UpsertChatDiffStatusParams) (database.ChatDiffStatus, error) {
|
||||
upsertCount.Add(1)
|
||||
return database.ChatDiffStatus{ChatID: arg.ChatID}, nil
|
||||
})
|
||||
|
||||
pub := func(_ context.Context, _ uuid.UUID) error {
|
||||
// Only the successful row publishes.
|
||||
publishCount.Add(1)
|
||||
signalIfDone()
|
||||
return nil
|
||||
}
|
||||
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
|
||||
// Fail ResolveBranchPullRequest for the first call, succeed
|
||||
// for the second.
|
||||
var callCount atomic.Int32
|
||||
refresher := newTestRefresher(t, mClock, withResolveBranchPR(
|
||||
func(context.Context, string, gitprovider.BranchRef) (*gitprovider.PRRef, error) {
|
||||
n := callCount.Add(1)
|
||||
if n == 1 {
|
||||
return nil, fmt.Errorf("simulated provider error")
|
||||
}
|
||||
return &gitprovider.PRRef{Owner: "o", Repo: "r", Number: 1}, nil
|
||||
},
|
||||
))
|
||||
|
||||
worker := gitsync.NewWorker(store, refresher, pub, mClock, logger)
|
||||
|
||||
tickOnce(ctx, t, mClock, worker, tickDone)
|
||||
|
||||
// BackoffChatDiffStatus was called for the failed row.
|
||||
assert.Equal(t, int32(1), backoffCount.Load())
|
||||
mu.Lock()
|
||||
require.Len(t, backoffArgs, 1)
|
||||
assert.Equal(t, chat1, backoffArgs[0].ChatID)
|
||||
// stale_at should be approximately clock.Now() + DiffStatusTTL (120s).
|
||||
expectedStaleAt := mClock.Now().UTC().Add(gitsync.DiffStatusTTL)
|
||||
assert.WithinDuration(t, expectedStaleAt, backoffArgs[0].StaleAt, time.Second)
|
||||
mu.Unlock()
|
||||
|
||||
// UpsertChatDiffStatus was called for the successful row.
|
||||
assert.Equal(t, int32(1), upsertCount.Load())
|
||||
// PublishDiffStatusChange was called only for the successful row.
|
||||
assert.Equal(t, int32(1), publishCount.Load())
|
||||
}
|
||||
|
||||
func TestWorker_UpsertError_ContinuesNextRow(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
chat1 := uuid.New()
|
||||
chat2 := uuid.New()
|
||||
ownerID := uuid.New()
|
||||
|
||||
var publishCount atomic.Int32
|
||||
tickDone := make(chan struct{})
|
||||
var closeOnce sync.Once
|
||||
var mu sync.Mutex
|
||||
upsertedChatIDs := make(map[uuid.UUID]struct{})
|
||||
|
||||
// We have 2 rows. The upsert for chat1 fails; the upsert
|
||||
// for chat2 succeeds and publishes. Because goroutines run
|
||||
// concurrently we don't know which finishes last, so we
|
||||
// track the total number of "terminal" events (upsert error
|
||||
// + publish success) and close tickDone when both have
|
||||
// occurred.
|
||||
var terminalOps atomic.Int32
|
||||
signalIfDone := func() {
|
||||
if terminalOps.Add(1) == 2 {
|
||||
closeOnce.Do(func() { close(tickDone) })
|
||||
}
|
||||
}
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().AcquireStaleChatDiffStatuses(gomock.Any(), gomock.Any()).
|
||||
Return([]database.AcquireStaleChatDiffStatusesRow{
|
||||
makeAcquiredRow(chat1, ownerID),
|
||||
makeAcquiredRow(chat2, ownerID),
|
||||
}, nil)
|
||||
store.EXPECT().UpsertChatDiffStatus(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, arg database.UpsertChatDiffStatusParams) (database.ChatDiffStatus, error) {
|
||||
if arg.ChatID == chat1 {
|
||||
// Terminal event for the failing row.
|
||||
signalIfDone()
|
||||
return database.ChatDiffStatus{}, fmt.Errorf("db write error")
|
||||
}
|
||||
mu.Lock()
|
||||
upsertedChatIDs[arg.ChatID] = struct{}{}
|
||||
mu.Unlock()
|
||||
return database.ChatDiffStatus{ChatID: arg.ChatID}, nil
|
||||
}).Times(2)
|
||||
|
||||
pub := func(_ context.Context, _ uuid.UUID) error {
|
||||
publishCount.Add(1)
|
||||
// Terminal event for the successful row.
|
||||
signalIfDone()
|
||||
return nil
|
||||
}
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
worker := gitsync.NewWorker(store, refresher, pub, mClock, logger)
|
||||
|
||||
tickOnce(ctx, t, mClock, worker, tickDone)
|
||||
|
||||
mu.Lock()
|
||||
_, gotChat2 := upsertedChatIDs[chat2]
|
||||
mu.Unlock()
|
||||
assert.True(t, gotChat2, "chat2 should have been upserted")
|
||||
assert.Equal(t, int32(1), publishCount.Load())
|
||||
}
|
||||
|
||||
func TestWorker_RespectsShutdown(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().AcquireStaleChatDiffStatuses(gomock.Any(), gomock.Any()).
|
||||
Return(nil, nil).AnyTimes()
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
worker := gitsync.NewWorker(store, refresher, nil, mClock, logger)
|
||||
|
||||
trap := mClock.Trap().NewTicker("gitsync", "worker")
|
||||
defer trap.Close()
|
||||
|
||||
workerCtx, cancel := context.WithCancel(ctx)
|
||||
go worker.Start(workerCtx)
|
||||
|
||||
// Wait for ticker creation so the worker is running.
|
||||
trap.MustWait(ctx).MustRelease(ctx)
|
||||
|
||||
// Cancel immediately.
|
||||
cancel()
|
||||
|
||||
select {
|
||||
case <-worker.Done():
|
||||
// Success — worker shut down.
|
||||
case <-ctx.Done():
|
||||
t.Fatal("timed out waiting for worker to shut down")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorker_MarkStale_UpsertAndPublish(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
workspaceID := uuid.New()
|
||||
ownerID := uuid.New()
|
||||
chat1 := uuid.New()
|
||||
chat2 := uuid.New()
|
||||
chatOther := uuid.New()
|
||||
|
||||
var mu sync.Mutex
|
||||
var upsertRefCalls []database.UpsertChatDiffStatusReferenceParams
|
||||
var publishedIDs []uuid.UUID
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().GetChatsByOwnerID(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, arg database.GetChatsByOwnerIDParams) ([]database.Chat, error) {
|
||||
require.Equal(t, ownerID, arg.OwnerID)
|
||||
return []database.Chat{
|
||||
{ID: chat1, OwnerID: ownerID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}},
|
||||
{ID: chat2, OwnerID: ownerID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}},
|
||||
{ID: chatOther, OwnerID: ownerID, WorkspaceID: uuid.NullUUID{UUID: uuid.New(), Valid: true}},
|
||||
}, nil
|
||||
})
|
||||
store.EXPECT().UpsertChatDiffStatusReference(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, arg database.UpsertChatDiffStatusReferenceParams) (database.ChatDiffStatus, error) {
|
||||
mu.Lock()
|
||||
upsertRefCalls = append(upsertRefCalls, arg)
|
||||
mu.Unlock()
|
||||
return database.ChatDiffStatus{ChatID: arg.ChatID}, nil
|
||||
}).Times(2)
|
||||
|
||||
pub := func(_ context.Context, chatID uuid.UUID) error {
|
||||
mu.Lock()
|
||||
publishedIDs = append(publishedIDs, chatID)
|
||||
mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
now := mClock.Now()
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
worker := gitsync.NewWorker(store, refresher, pub, mClock, logger)
|
||||
|
||||
worker.MarkStale(ctx, workspaceID, ownerID, "feature", "https://github.com/owner/repo")
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
require.Len(t, upsertRefCalls, 2)
|
||||
for _, call := range upsertRefCalls {
|
||||
assert.Equal(t, "feature", call.GitBranch)
|
||||
assert.Equal(t, "https://github.com/owner/repo", call.GitRemoteOrigin)
|
||||
assert.True(t, call.StaleAt.Before(now),
|
||||
"stale_at should be in the past, got %v vs now %v", call.StaleAt, now)
|
||||
assert.Equal(t, sql.NullString{}, call.Url)
|
||||
}
|
||||
|
||||
require.Len(t, publishedIDs, 2)
|
||||
assert.ElementsMatch(t, []uuid.UUID{chat1, chat2}, publishedIDs)
|
||||
}
|
||||
|
||||
func TestWorker_MarkStale_NoMatchingChats(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
workspaceID := uuid.New()
|
||||
ownerID := uuid.New()
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().GetChatsByOwnerID(gomock.Any(), gomock.Any()).
|
||||
Return([]database.Chat{
|
||||
{ID: uuid.New(), OwnerID: ownerID, WorkspaceID: uuid.NullUUID{UUID: uuid.New(), Valid: true}},
|
||||
{ID: uuid.New(), OwnerID: ownerID, WorkspaceID: uuid.NullUUID{UUID: uuid.New(), Valid: true}},
|
||||
}, nil)
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
worker := gitsync.NewWorker(store, refresher, nil, mClock, logger)
|
||||
|
||||
worker.MarkStale(ctx, workspaceID, ownerID, "main", "https://github.com/x/y")
|
||||
}
|
||||
|
||||
func TestWorker_MarkStale_UpsertFails_ContinuesNext(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
workspaceID := uuid.New()
|
||||
ownerID := uuid.New()
|
||||
chat1 := uuid.New()
|
||||
chat2 := uuid.New()
|
||||
|
||||
var publishCount atomic.Int32
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().GetChatsByOwnerID(gomock.Any(), gomock.Any()).
|
||||
Return([]database.Chat{
|
||||
{ID: chat1, OwnerID: ownerID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}},
|
||||
{ID: chat2, OwnerID: ownerID, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}},
|
||||
}, nil)
|
||||
store.EXPECT().UpsertChatDiffStatusReference(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, arg database.UpsertChatDiffStatusReferenceParams) (database.ChatDiffStatus, error) {
|
||||
if arg.ChatID == chat1 {
|
||||
return database.ChatDiffStatus{}, fmt.Errorf("upsert ref error")
|
||||
}
|
||||
return database.ChatDiffStatus{ChatID: arg.ChatID}, nil
|
||||
}).Times(2)
|
||||
|
||||
pub := func(_ context.Context, _ uuid.UUID) error {
|
||||
publishCount.Add(1)
|
||||
return nil
|
||||
}
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
worker := gitsync.NewWorker(store, refresher, pub, mClock, logger)
|
||||
|
||||
worker.MarkStale(ctx, workspaceID, ownerID, "dev", "https://github.com/a/b")
|
||||
|
||||
assert.Equal(t, int32(1), publishCount.Load())
|
||||
}
|
||||
|
||||
func TestWorker_MarkStale_GetChatsFails(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().GetChatsByOwnerID(gomock.Any(), gomock.Any()).
|
||||
Return(nil, fmt.Errorf("db error"))
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
worker := gitsync.NewWorker(store, refresher, nil, mClock, logger)
|
||||
|
||||
worker.MarkStale(ctx, uuid.New(), uuid.New(), "main", "https://github.com/x/y")
|
||||
}
|
||||
|
||||
func TestWorker_TickStoreError(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
tickDone := make(chan struct{})
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
store.EXPECT().AcquireStaleChatDiffStatuses(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(context.Context, int32) ([]database.AcquireStaleChatDiffStatusesRow, error) {
|
||||
close(tickDone)
|
||||
return nil, fmt.Errorf("database unavailable")
|
||||
})
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
worker := gitsync.NewWorker(store, refresher, nil, mClock, logger)
|
||||
|
||||
tickOnce(ctx, t, mClock, worker, tickDone)
|
||||
}
|
||||
|
||||
func TestWorker_MarkStale_EmptyBranchOrOrigin(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
branch string
|
||||
origin string
|
||||
}{
|
||||
{"both empty", "", ""},
|
||||
{"branch empty", "", "https://github.com/x/y"},
|
||||
{"origin empty", "main", ""},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
store := dbmock.NewMockStore(ctrl)
|
||||
|
||||
mClock := quartz.NewMock(t)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
worker := gitsync.NewWorker(store, refresher, nil, mClock, logger)
|
||||
|
||||
worker.MarkStale(ctx, uuid.New(), uuid.New(), tc.branch, tc.origin)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestWorker exercises the worker tick against a
|
||||
// real PostgreSQL database to verify that the SQL queries, foreign key
|
||||
// constraints, and upsert logic work end-to-end.
|
||||
func TestWorker(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
// 1. Real database store.
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
|
||||
// 2. Create a user (FK for chats).
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
|
||||
// 3. Set up FK chain: chat_providers -> chat_model_configs -> chats.
|
||||
_, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{
|
||||
Provider: "openai",
|
||||
DisplayName: "OpenAI",
|
||||
Enabled: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
modelCfg, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
|
||||
Provider: "openai",
|
||||
Model: "test-model",
|
||||
DisplayName: "Test Model",
|
||||
Enabled: true,
|
||||
ContextLimit: 100000,
|
||||
CompressionThreshold: 70,
|
||||
Options: json.RawMessage("{}"),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
chat, err := db.InsertChat(ctx, database.InsertChatParams{
|
||||
OwnerID: user.ID,
|
||||
LastModelConfigID: modelCfg.ID,
|
||||
Title: "integration-test",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// 4. Seed a stale diff status row so the worker picks it up.
|
||||
_, err = db.UpsertChatDiffStatusReference(ctx, database.UpsertChatDiffStatusReferenceParams{
|
||||
ChatID: chat.ID,
|
||||
GitBranch: "feature",
|
||||
GitRemoteOrigin: "https://github.com/o/r",
|
||||
StaleAt: time.Now().Add(-time.Minute),
|
||||
Url: sql.NullString{},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// 5. Mock refresher returns a canned PR status.
|
||||
mClock := quartz.NewMock(t)
|
||||
refresher := newTestRefresher(t, mClock)
|
||||
|
||||
// 6. Track publish calls.
|
||||
var publishCount atomic.Int32
|
||||
tickDone := make(chan struct{})
|
||||
pub := func(_ context.Context, chatID uuid.UUID) error {
|
||||
assert.Equal(t, chat.ID, chatID)
|
||||
if publishCount.Add(1) == 1 {
|
||||
close(tickDone)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 7. Create and run the worker for one tick.
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
worker := gitsync.NewWorker(db, refresher, pub, mClock, logger)
|
||||
|
||||
tickOnce(ctx, t, mClock, worker, tickDone)
|
||||
|
||||
// 8. Assert publisher was called.
|
||||
require.Equal(t, int32(1), publishCount.Load())
|
||||
|
||||
// 9. Read back and verify persisted fields.
|
||||
status, err := db.GetChatDiffStatusByChatID(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// The mock resolveBranchPR returns PRRef{Owner: "o", Repo: "r", Number: 1}
|
||||
// and buildPullRequestURL formats it as https://github.com/o/r/pull/1.
|
||||
assert.Equal(t, "https://github.com/o/r/pull/1", status.Url.String)
|
||||
assert.True(t, status.Url.Valid)
|
||||
assert.Equal(t, string(gitprovider.PRStateOpen), status.PullRequestState.String)
|
||||
assert.True(t, status.PullRequestState.Valid)
|
||||
assert.Equal(t, int32(10), status.Additions)
|
||||
assert.Equal(t, int32(3), status.Deletions)
|
||||
assert.Equal(t, int32(2), status.ChangedFiles)
|
||||
assert.True(t, status.RefreshedAt.Valid, "refreshed_at should be set")
|
||||
// The mock clock's Now() + DiffStatusTTL determines stale_at.
|
||||
expectedStaleAt := mClock.Now().Add(gitsync.DiffStatusTTL)
|
||||
assert.WithinDuration(t, expectedStaleAt, status.StaleAt, time.Second)
|
||||
}
|
||||
Reference in New Issue
Block a user