coder · deansheather · Feb 12, 2024 · Feb 12, 2024
diff --git a/agent/agenttest/client.go b/agent/agenttest/client.go
@@ -197,9 +197,10 @@ type FakeAgentAPI struct {
 	t      testing.TB
 	logger slog.Logger
 
-	manifest  *agentproto.Manifest
-	startupCh chan *agentproto.Startup
-	statsCh   chan *agentproto.Stats
+	manifest    *agentproto.Manifest
+	startupCh   chan *agentproto.Startup
+	statsCh     chan *agentproto.Stats
+	appHealthCh chan *agentproto.BatchUpdateAppHealthRequest
 
 	getServiceBannerFunc func() (codersdk.ServiceBannerConfig, error)
 }
@@ -244,9 +245,14 @@ func (*FakeAgentAPI) UpdateLifecycle(context.Context, *agentproto.UpdateLifecycl
 
 func (f *FakeAgentAPI) BatchUpdateAppHealths(ctx context.Context, req *agentproto.BatchUpdateAppHealthRequest) (*agentproto.BatchUpdateAppHealthResponse, error) {
 	f.logger.Debug(ctx, "batch update app health", slog.F("req", req))
+	f.appHealthCh <- req
 	return &agentproto.BatchUpdateAppHealthResponse{}, nil
 }
 
+func (f *FakeAgentAPI) AppHealthCh() <-chan *agentproto.BatchUpdateAppHealthRequest {
+	return f.appHealthCh
+}
+
 func (f *FakeAgentAPI) UpdateStartup(_ context.Context, req *agentproto.UpdateStartupRequest) (*agentproto.Startup, error) {
 	f.startupCh <- req.GetStartup()
 	return req.GetStartup(), nil
@@ -264,10 +270,11 @@ func (*FakeAgentAPI) BatchCreateLogs(context.Context, *agentproto.BatchCreateLog
 
 func NewFakeAgentAPI(t testing.TB, logger slog.Logger, manifest *agentproto.Manifest, statsCh chan *agentproto.Stats) *FakeAgentAPI {
 	return &FakeAgentAPI{
-		t:         t,
-		logger:    logger.Named("FakeAgentAPI"),
-		manifest:  manifest,
-		statsCh:   statsCh,
-		startupCh: make(chan *agentproto.Startup, 100),
+		t:           t,
+		logger:      logger.Named("FakeAgentAPI"),
+		manifest:    manifest,
+		statsCh:     statsCh,
+		startupCh:   make(chan *agentproto.Startup, 100),
+		appHealthCh: make(chan *agentproto.BatchUpdateAppHealthRequest, 100),
 	}
 }
diff --git a/agent/apphealth_test.go b/agent/apphealth_test.go
@@ -4,16 +4,21 @@ import (
 	"context"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/v2/agent"
+	"github.com/coder/coder/v2/agent/agenttest"
+	"github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -40,12 +45,23 @@ func TestAppHealth_Healthy(t *testing.T) {
 			},
 			Health: codersdk.WorkspaceAppHealthInitializing,
 		},
+		{
+			Slug: "app3",
+			Healthcheck: codersdk.Healthcheck{
+				Interval:  2,
+				Threshold: 1,
+			},
+			Health: codersdk.WorkspaceAppHealthInitializing,
+		},
 	}
 	handlers := []http.Handler{
 		nil,
 		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			httpapi.Write(r.Context(), w, http.StatusOK, nil)
 		}),
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			httpapi.Write(r.Context(), w, http.StatusOK, nil)
+		}),
 	}
 	getApps, closeFn := setupAppReporter(ctx, t, apps, handlers)
 	defer closeFn()
@@ -58,7 +74,7 @@ func TestAppHealth_Healthy(t *testing.T) {
 			return false
 		}
 
-		return apps[1].Health == codersdk.WorkspaceAppHealthHealthy
+		return apps[1].Health == codersdk.WorkspaceAppHealthHealthy && apps[2].Health == codersdk.WorkspaceAppHealthHealthy
 	}, testutil.WaitLong, testutil.IntervalSlow)
 }
 
@@ -163,6 +179,12 @@ func TestAppHealth_NotSpamming(t *testing.T) {
 
 func setupAppReporter(ctx context.Context, t *testing.T, apps []codersdk.WorkspaceApp, handlers []http.Handler) (agent.WorkspaceAgentApps, func()) {
 	closers := []func(){}
+	for i, app := range apps {
+		if app.ID == uuid.Nil {
+			app.ID = uuid.New()
+			apps[i] = app
+		}
+	}
 	for i, handler := range handlers {
 		if handler == nil {
 			continue
@@ -181,23 +203,43 @@ func setupAppReporter(ctx context.Context, t *testing.T, apps []codersdk.Workspa
 		var newApps []codersdk.WorkspaceApp
 		return append(newApps, apps...), nil
 	}
-	postWorkspaceAgentAppHealth := func(_ context.Context, req agentsdk.PostAppHealthsRequest) error {
-		mu.Lock()
-		for id, health := range req.Healths {
-			for i, app := range apps {
-				if app.ID != id {
-					continue
+
+	// We don't care about manifest or stats in this test since it's not using
+	// a full agent and these RPCs won't get called.
+	//
+	// We use a proper fake agent API so we can test the conversion code and the
+	// request code as well. Before we were bypassing these by using a custom
+	// post function.
+	fakeAAPI := agenttest.NewFakeAgentAPI(t, slogtest.Make(t, nil), nil, nil)
+
+	// Process events from the channel and update the health of the apps.
+	go func() {
+		appHealthCh := fakeAAPI.AppHealthCh()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case req := <-appHealthCh:
+				mu.Lock()
+				for _, update := range req.Updates {
+					updateID, err := uuid.FromBytes(update.Id)
+					assert.NoError(t, err)
+					updateHealth := codersdk.WorkspaceAppHealth(strings.ToLower(proto.AppHealth_name[int32(update.Health)]))
+
+					for i, app := range apps {
+						if app.ID != updateID {
+							continue
+						}
+						app.Health = updateHealth
+						apps[i] = app
+					}
 				}
-				app.Health = health
-				apps[i] = app
+				mu.Unlock()
 			}
 		}
-		mu.Unlock()
-
-		return nil
-	}
+	}()
 
-	go agent.NewWorkspaceAppHealthReporter(slogtest.Make(t, nil).Leveled(slog.LevelDebug), apps, postWorkspaceAgentAppHealth)(ctx)
+	go agent.NewWorkspaceAppHealthReporter(slogtest.Make(t, nil).Leveled(slog.LevelDebug), apps, agentsdk.AppHealthPoster(fakeAAPI))(ctx)
 
 	return workspaceAgentApps, func() {
 		for _, closeFn := range closers {

diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -226,7 +226,12 @@ type PostAppHealthsRequest struct {
 	Healths map[uuid.UUID]codersdk.WorkspaceAppHealth
 }
 
-func AppHealthPoster(aAPI proto.DRPCAgentClient) func(ctx context.Context, req PostAppHealthsRequest) error {
+// BatchUpdateAppHealthsClient is a partial interface of proto.DRPCAgentClient.
+type BatchUpdateAppHealthsClient interface {
+	BatchUpdateAppHealths(ctx context.Context, req *proto.BatchUpdateAppHealthRequest) (*proto.BatchUpdateAppHealthResponse, error)
+}
+
+func AppHealthPoster(aAPI BatchUpdateAppHealthsClient) func(ctx context.Context, req PostAppHealthsRequest) error {
 	return func(ctx context.Context, req PostAppHealthsRequest) error {
 		pReq, err := ProtoFromAppHealthsRequest(req)
 		if err != nil {

diff --git a/codersdk/agentsdk/convert.go b/codersdk/agentsdk/convert.go
@@ -287,6 +287,8 @@ func ProtoFromAppHealthsRequest(req PostAppHealthsRequest) (*proto.BatchUpdateAp
 			return nil, xerrors.Errorf("unknown app health: %s", h)
 		}
 
+		// Copy the ID, otherwise all updates will have the same ID (the last
+		// one in the list).
 		var idCopy uuid.UUID
 		copy(idCopy[:], id[:])
 		pReq.Updates = append(pReq.Updates, &proto.BatchUpdateAppHealthRequest_HealthUpdate{