chore: ensure we close and init

DanielleMaywood · DanielleMaywood · commit ae5dd1e08690 · 2025-06-24T22:43:46.000Z
diff --git a/agent/agentcontainers/api.go b/agent/agentcontainers/api.go
@@ -265,8 +265,6 @@ func NewAPI(logger slog.Logger, options ...Option) *API {
 	api := &API{
 		ctx:                         ctx,
 		cancel:                      cancel,
-		watcherDone:                 make(chan struct{}),
-		updaterDone:                 make(chan struct{}),
 		initialUpdateDone:           make(chan struct{}),
 		updateTrigger:               make(chan chan error),
 		updateInterval:              defaultUpdateInterval,
@@ -326,6 +324,9 @@ func (api *API) Init(opts ...Option) {
 		opt(api)
 	}
 
+	api.watcherDone = make(chan struct{})
+	api.updaterDone = make(chan struct{})
+
 	go api.watcherLoop()
 	go api.updaterLoop()
 }
@@ -1640,8 +1641,12 @@ func (api *API) Close() error {
 	err := api.watcher.Close()
 
 	// Wait for loops to finish.
-	<-api.watcherDone
-	<-api.updaterDone
+	if api.watcherDone != nil {
+		<-api.watcherDone
+	}
+	if api.updaterDone != nil {
+		<-api.updaterDone
+	}
 
 	// Wait for all async tasks to complete.
 	api.asyncWg.Wait()
diff --git a/agent/agentcontainers/api_test.go b/agent/agentcontainers/api_test.go
@@ -437,6 +437,7 @@ func TestAPI(t *testing.T) {
 					agentcontainers.WithContainerCLI(mLister),
 					agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
 				)
+				api.Init()
 				defer api.Close()
 				r.Mount("/", api.Routes())
 
@@ -614,6 +615,7 @@ func TestAPI(t *testing.T) {
 					agentcontainers.WithDevcontainerCLI(tt.devcontainerCLI),
 					agentcontainers.WithWatcher(watcher.NewNoop()),
 				)
+				api.Init()
 				defer api.Close()
 				r.Mount("/", api.Routes())
 
@@ -1027,6 +1029,7 @@ func TestAPI(t *testing.T) {
 				}
 
 				api := agentcontainers.NewAPI(logger, apiOptions...)
+				api.Init()
 				defer api.Close()
 
 				r.Mount("/", api.Routes())
@@ -1111,6 +1114,7 @@ func TestAPI(t *testing.T) {
 				[]codersdk.WorkspaceAgentScript{{LogSourceID: uuid.New(), ID: dc.ID}},
 			),
 		)
+		api.Init()
 		defer api.Close()
 
 		// Make sure the ticker function has been registered
@@ -1206,6 +1210,7 @@ func TestAPI(t *testing.T) {
 			agentcontainers.WithWatcher(fWatcher),
 			agentcontainers.WithClock(mClock),
 		)
+		api.Init()
 		defer api.Close()
 
 		r := chi.NewRouter()
@@ -1358,6 +1363,7 @@ func TestAPI(t *testing.T) {
 			agentcontainers.WithDevcontainerCLI(fakeDCCLI),
 			agentcontainers.WithManifestInfo("test-user", "test-workspace"),
 		)
+		api.Init()
 		apiClose := func() {
 			closeOnce.Do(func() {
 				// Close before api.Close() defer to avoid deadlock after test.
@@ -1578,6 +1584,7 @@ func TestAPI(t *testing.T) {
 			agentcontainers.WithSubAgentClient(fakeSAC),
 			agentcontainers.WithDevcontainerCLI(&fakeDevcontainerCLI{}),
 		)
+		api.Init()
 		defer api.Close()
 
 		tickerTrap.MustWait(ctx).MustRelease(ctx)
@@ -1899,6 +1906,7 @@ func TestAPI(t *testing.T) {
 					agentcontainers.WithSubAgentURL("test-subagent-url"),
 					agentcontainers.WithWatcher(watcher.NewNoop()),
 				)
+				api.Init()
 				defer api.Close()
 
 				// Close before api.Close() defer to avoid deadlock after test.
@@ -1991,6 +1999,7 @@ func TestAPI(t *testing.T) {
 			agentcontainers.WithSubAgentURL("test-subagent-url"),
 			agentcontainers.WithWatcher(watcher.NewNoop()),
 		)
+		api.Init()
 		defer api.Close()
 
 		// Close before api.Close() defer to avoid deadlock after test.
@@ -2045,6 +2054,7 @@ func TestAPI(t *testing.T) {
 			agentcontainers.WithExecer(fakeExec),
 			agentcontainers.WithCommandEnv(commandEnv),
 		)
+		api.Init()
 		defer api.Close()
 
 		// Call RefreshContainers directly to trigger CommandEnv usage.
@@ -2134,6 +2144,7 @@ func TestAPI(t *testing.T) {
 			agentcontainers.WithWatcher(fWatcher),
 			agentcontainers.WithClock(mClock),
 		)
+		api.Init()
 		defer func() {
 			close(fakeSAC.createErrC)
 			close(fakeSAC.deleteErrC)
@@ -2334,6 +2345,7 @@ func TestSubAgentCreationWithNameRetry(t *testing.T) {
 				agentcontainers.WithSubAgentClient(fSAC),
 				agentcontainers.WithWatcher(watcher.NewNoop()),
 			)
+			api.Init()
 			defer api.Close()
 
 			tickerTrap.MustWait(ctx).MustRelease(ctx)
