fix: send end of logs when dbfake completes job (#11402)

This commit is contained in:
Spike Curtis 2024-01-04 15:18:00 +04:00 committed by GitHub
parent 48cd4c3a10
commit f9ebe8c719
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 19 additions and 3 deletions

View File

@ -12,6 +12,7 @@ import (
"github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbfake" "github.com/coder/coder/v2/coderd/database/dbfake"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/provisioner/echo" "github.com/coder/coder/v2/provisioner/echo"
"github.com/coder/coder/v2/provisionersdk/proto" "github.com/coder/coder/v2/provisionersdk/proto"
@ -412,10 +413,11 @@ func TestStart_Starting(t *testing.T) {
t.Parallel() t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort) ctx := testutil.Context(t, testutil.WaitShort)
client, db := coderdtest.NewWithDatabase(t, nil) store, ps := dbtestutil.NewDB(t)
client := coderdtest.New(t, &coderdtest.Options{Pubsub: ps, Database: store})
owner := coderdtest.CreateFirstUser(t, client) owner := coderdtest.CreateFirstUser(t, client)
memberClient, member := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID) memberClient, member := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
r := dbfake.WorkspaceBuild(t, db, database.Workspace{ r := dbfake.WorkspaceBuild(t, store, database.Workspace{
OwnerID: member.ID, OwnerID: member.ID,
OrganizationID: owner.OrganizationID, OrganizationID: owner.OrganizationID,
}). }).
@ -434,7 +436,7 @@ func TestStart_Starting(t *testing.T) {
pty.ExpectMatch("workspace is already starting") pty.ExpectMatch("workspace is already starting")
_ = dbfake.JobComplete(t, db, r.Build.JobID).Do() _ = dbfake.JobComplete(t, store, r.Build.JobID).Pubsub(ps).Do()
pty.ExpectMatch("workspace has been started") pty.ExpectMatch("workspace has been started")
_ = testutil.RequireRecvCtx(ctx, t, doneChan) _ = testutil.RequireRecvCtx(ctx, t, doneChan)

View File

@ -20,6 +20,7 @@ import (
"github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/coderd/telemetry" "github.com/coder/coder/v2/coderd/telemetry"
"github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/provisionersdk"
sdkproto "github.com/coder/coder/v2/provisionersdk/proto" sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
) )
@ -379,6 +380,7 @@ type JobCompleteBuilder struct {
t testing.TB t testing.TB
db database.Store db database.Store
jobID uuid.UUID jobID uuid.UUID
ps pubsub.Pubsub
} }
type JobCompleteResponse struct { type JobCompleteResponse struct {
@ -393,6 +395,12 @@ func JobComplete(t testing.TB, db database.Store, jobID uuid.UUID) JobCompleteBu
} }
} }
func (b JobCompleteBuilder) Pubsub(ps pubsub.Pubsub) JobCompleteBuilder {
// nolint: revive // returns modified struct
b.ps = ps
return b
}
func (b JobCompleteBuilder) Do() JobCompleteResponse { func (b JobCompleteBuilder) Do() JobCompleteResponse {
r := JobCompleteResponse{CompletedAt: dbtime.Now()} r := JobCompleteResponse{CompletedAt: dbtime.Now()}
err := b.db.UpdateProvisionerJobWithCompleteByID(ownerCtx, database.UpdateProvisionerJobWithCompleteByIDParams{ err := b.db.UpdateProvisionerJobWithCompleteByID(ownerCtx, database.UpdateProvisionerJobWithCompleteByIDParams{
@ -406,6 +414,12 @@ func (b JobCompleteBuilder) Do() JobCompleteResponse {
}, },
}) })
require.NoError(b.t, err, "complete job") require.NoError(b.t, err, "complete job")
if b.ps != nil {
data, err := json.Marshal(provisionersdk.ProvisionerJobLogsNotifyMessage{EndOfLogs: true})
require.NoError(b.t, err)
err = b.ps.Publish(provisionersdk.ProvisionerJobLogsNotifyChannel(b.jobID), data)
require.NoError(b.t, err)
}
return r return r
} }