add flush to wsproxy tests

johnstcn · johnstcn · commit 601136cd23b4 · 2024-01-16T13:48:43.000Z
diff --git a/enterprise/coderd/coderdenttest/proxytest.go b/enterprise/coderd/coderdenttest/proxytest.go
@@ -37,6 +37,9 @@ type ProxyOptions struct {
 
 	// ProxyURL is optional
 	ProxyURL *url.URL
+
+	// Flush is optional
+	Flush chan chan<- struct{}
 }
 
 // NewWorkspaceProxy will configure a wsproxy.Server with the given options.
@@ -113,6 +116,9 @@ func NewWorkspaceProxy(t *testing.T, coderdAPI *coderd.API, owner *codersdk.Clie
 	// Inherit collector options from coderd, but keep the wsproxy reporter.
 	statsCollectorOptions := coderdAPI.Options.WorkspaceAppsStatsCollectorOptions
 	statsCollectorOptions.Reporter = nil
+	if options.Flush != nil {
+		statsCollectorOptions.Flush = options.Flush
+	}
 
 	wssrv, err := wsproxy.New(ctx, &wsproxy.Options{
 		Logger:            slogtest.Make(t, nil).Leveled(slog.LevelDebug),
diff --git a/enterprise/wsproxy/wsproxy_test.go b/enterprise/wsproxy/wsproxy_test.go
@@ -442,6 +442,13 @@ func TestWorkspaceProxyWorkspaceApps(t *testing.T) {
 			"*",
 		}
 
+		proxyFlushCh := make(chan chan<- struct{}, 1)
+		flushStats := func() {
+			proxyFlushDone := make(chan struct{}, 1)
+			proxyFlushCh <- proxyFlushDone
+			<-proxyFlushDone
+		}
+
 		client, closer, api, user := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
 			Options: &coderdtest.Options{
 				DeploymentValues:         deploymentValues,
@@ -476,13 +483,15 @@ func TestWorkspaceProxyWorkspaceApps(t *testing.T) {
 			Name:            "best-proxy",
 			AppHostname:     opts.AppHost,
 			DisablePathApps: opts.DisablePathApps,
+			Flush:           proxyFlushCh,
 		})
 
 		return &apptest.Deployment{
 			Options:        opts,
 			SDKClient:      client,
 			FirstUser:      user,
 			PathAppBaseURL: proxyAPI.Options.AccessURL,
+			FlushStats:     flushStats,
 		}
 	})
 }
