fixup! refactor(container): restart handling

internal/container/container.go

@@ -377,12 +377,7 @@ func (a *Client) runContainerLoop(
 			respC, errC := a.apiClient.ContainerWait(ctx, containerID, container.WaitConditionNextExit)
 			select {
 			case resp := <-respC:
-				if shouldRestartFunc != nil {
-					shouldRestart, err := shouldRestartFunc(resp.StatusCode, restartCount, time.Since(startedWaitingAt))
-					if shouldRestart && err != nil {
-						panic(fmt.Errorf("shouldRestart must return nil error if restarting, but returned: %w", err))
-					}
-					if !shouldRestart {
+				exit := func(err error) {
 					a.logger.Info("Container exited", "id", shortID(containerID), "should_restart", "false", "exit_code", resp.StatusCode, "restart_count", restartCount)
 					containerRespC <- containerWaitResponse{
 						WaitResponse: resp,
@@ -390,8 +385,20 @@ func (a *Client) runContainerLoop(
 						restartCount: restartCount,
 						err:          err,
 					}
+				}
+
+				if shouldRestartFunc == nil {
+					exit(nil)
 					return
 				}
+
+				shouldRestart, err := shouldRestartFunc(resp.StatusCode, restartCount, time.Since(startedWaitingAt))
+				if shouldRestart && err != nil {
+					panic(fmt.Errorf("shouldRestart must return nil error if restarting, but returned: %w", err))
+				}
+				if !shouldRestart {
+					exit(err)
+					return
 				}
 
 				a.logger.Info("Container exited", "id", shortID(containerID), "should_restart", "true", "exit_code", resp.StatusCode, "restart_count", restartCount)

internal/container/container_test.go

@@ -94,9 +94,6 @@ func TestClientRunContainer(t *testing.T) {
 				Mode:    0755,
 			},
 		},
-		ShouldRestart: func(int64, int, time.Duration) (bool, error) {
-			return false, nil
-		},
 	})
 
 	done := make(chan struct{})