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

spikecurtis · web-flow · commit f9ebe8c71909 · 2024-01-04T15:18:00.000+04:00
diff --git a/cli/start_test.go b/cli/start_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"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/provisioner/echo"
 	"github.com/coder/coder/v2/provisionersdk/proto"
@@ -412,10 +413,11 @@ func TestStart_Starting(t *testing.T) {
 	t.Parallel()
 	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)
 	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,
 		OrganizationID: owner.OrganizationID,
 	}).
@@ -434,7 +436,7 @@ func TestStart_Starting(t *testing.T) {
 
 	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")
 
 	_ = testutil.RequireRecvCtx(ctx, t, doneChan)
diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -20,6 +20,7 @@ import (
 	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/coderd/telemetry"
 	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/provisionersdk"
 	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
 )
 
@@ -379,6 +380,7 @@ type JobCompleteBuilder struct {
 	t     testing.TB
 	db    database.Store
 	jobID uuid.UUID
+	ps    pubsub.Pubsub
 }
 
 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 {
 	r := JobCompleteResponse{CompletedAt: dbtime.Now()}
 	err := b.db.UpdateProvisionerJobWithCompleteByID(ownerCtx, database.UpdateProvisionerJobWithCompleteByIDParams{
@@ -406,6 +414,12 @@ func (b JobCompleteBuilder) Do() JobCompleteResponse {
 		},
 	})
 	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
 }
 
