add stats collection test to apptest

mafredri · mafredri · commit 97ef37b2446d · 2023-08-14T12:52:00.000Z
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -162,6 +162,8 @@ type Options struct {
 
 	UpdateAgentMetrics func(ctx context.Context, username, workspaceName, agentName string, metrics []agentsdk.AgentMetric)
 	StatsBatcher       *batchstats.Batcher
+
+	WorkspaceAppsStatsCollector *workspaceapps.StatsCollector
 }
 
 // @title Coder API
@@ -416,8 +418,13 @@ func New(options *Options) *API {
 			Cache: wsconncache.New(api._dialWorkspaceAgentTailnet, 0),
 		}
 	}
-
 	workspaceAppsLogger := options.Logger.Named("workspaceapps")
+	if options.WorkspaceAppsStatsCollector == nil {
+		options.WorkspaceAppsStatsCollector = workspaceapps.NewStatsCollector(workspaceapps.StatsCollectorOptions{
+			Logger:   workspaceAppsLogger.Named("stats_collector"),
+			Reporter: workspaceapps.NewStatsDBReporter(options.Database, workspaceapps.DefaultStatsDBReporterBatchSize),
+		})
+	}
 	api.workspaceAppServer = &workspaceapps.Server{
 		Logger: workspaceAppsLogger,
 
@@ -430,10 +437,7 @@ func New(options *Options) *API {
 		SignedTokenProvider: api.WorkspaceAppsProvider,
 		AgentProvider:       api.agentProvider,
 		AppSecurityKey:      options.AppSecurityKey,
-		StatsCollector: workspaceapps.NewStatsCollector(workspaceapps.StatsCollectorOptions{
-			Logger:   workspaceAppsLogger.Named("stats_collector"),
-			Reporter: workspaceapps.NewStatsDBReporter(options.Database, workspaceapps.DefaultStatsDBReporterBatchSize),
-		}),
+		StatsCollector:      options.WorkspaceAppsStatsCollector,
 
 		DisablePathApps:  options.DeploymentValues.DisablePathApps.Value(),
 		SecureAuthCookie: options.DeploymentValues.SecureAuthCookie.Value(),
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -144,6 +144,8 @@ type Options struct {
 	// as part of your test.
 	Logger       *slog.Logger
 	StatsBatcher *batchstats.Batcher
+
+	WorkspaceAppsStatsCollector *workspaceapps.StatsCollector
 }
 
 // New constructs a codersdk client connected to an in-memory API instance.
@@ -425,6 +427,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 			HealthcheckTimeout:          options.HealthcheckTimeout,
 			HealthcheckRefresh:          options.HealthcheckRefresh,
 			StatsBatcher:                options.StatsBatcher,
+			WorkspaceAppsStatsCollector: options.WorkspaceAppsStatsCollector,
 		}
 }
 
diff --git a/coderd/workspaceapps/apptest/apptest.go b/coderd/workspaceapps/apptest/apptest.go
@@ -15,6 +15,7 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -1406,4 +1407,51 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
 		require.Equal(t, []string{"Origin", "X-Foobar"}, deduped)
 		require.Equal(t, []string{"baz"}, resp.Header.Values("X-Foobar"))
 	})
+
+	t.Run("ReportStats", func(t *testing.T) {
+		t.Parallel()
+
+		reporter := &fakeStatsReporter{}
+		collector := workspaceapps.NewStatsCollector(workspaceapps.StatsCollectorOptions{
+			Reporter:       reporter,
+			ReportInterval: time.Second,
+			RollupWindow:   time.Minute,
+		})
+		appDetails := setupProxyTest(t, &DeploymentOptions{
+			StatsCollector: collector,
+		})
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		u := appDetails.PathAppURL(appDetails.Apps.Owner)
+		resp, err := requestWithRetries(ctx, t, appDetails.AppClient(t), http.MethodGet, u.String(), nil)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+		_, err = io.Copy(io.Discard, resp.Body)
+		require.NoError(t, err)
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+
+		require.Eventually(t, func() bool {
+			return len(reporter.stats()) > 0
+		}, testutil.WaitLong, testutil.IntervalMedium, "stats not reported")
+	})
+}
+
+type fakeStatsReporter struct {
+	mu sync.Mutex
+	s  []workspaceapps.StatsReport
+}
+
+func (r *fakeStatsReporter) stats() []workspaceapps.StatsReport {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.s
+}
+
+func (r *fakeStatsReporter) Report(_ context.Context, stats []workspaceapps.StatsReport) error {
+	r.mu.Lock()
+	r.s = append(r.s, stats...)
+	r.mu.Unlock()
+	return nil
 }
diff --git a/coderd/workspaceapps/apptest/setup.go b/coderd/workspaceapps/apptest/setup.go
@@ -21,6 +21,7 @@ import (
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/agent"
 	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/coderd/workspaceapps"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/codersdk/agentsdk"
 	"github.com/coder/coder/provisioner/echo"
@@ -51,6 +52,8 @@ type DeploymentOptions struct {
 	DangerousAllowPathAppSiteOwnerAccess bool
 	ServeHTTPS                           bool
 
+	StatsCollector *workspaceapps.StatsCollector
+
 	// The following fields are only used by setupProxyTestWithFactory.
 	noWorkspace bool
 	port        uint16
diff --git a/coderd/workspaceapps_test.go b/coderd/workspaceapps_test.go
@@ -275,6 +275,7 @@ func TestWorkspaceApps(t *testing.T) {
 					"CF-Connecting-IP",
 				},
 			},
+			WorkspaceAppsStatsCollector: opts.StatsCollector,
 		})
 
 		user := coderdtest.CreateFirstUser(t, client)
diff --git a/enterprise/wsproxy/wsproxy_test.go b/enterprise/wsproxy/wsproxy_test.go
@@ -596,6 +596,7 @@ func TestWorkspaceProxyWorkspaceApps_Wsconncache(t *testing.T) {
 						"CF-Connecting-IP",
 					},
 				},
+				WorkspaceAppsStatsCollector: opts.StatsCollector,
 			},
 			LicenseOptions: &coderdenttest.LicenseOptions{
 				Features: license.Features{
@@ -654,6 +655,7 @@ func TestWorkspaceProxyWorkspaceApps_SingleTailnet(t *testing.T) {
 						"CF-Connecting-IP",
 					},
 				},
+				WorkspaceAppsStatsCollector: opts.StatsCollector,
 			},
 			LicenseOptions: &coderdenttest.LicenseOptions{
 				Features: license.Features{

Original file line number	Diff line number	Diff line change
`@@ -144,6 +144,8 @@ type Options struct {`
`144`	`144`	`// as part of your test.`
`145`	`145`	`Logger *slog.Logger`
`146`	`146`	`StatsBatcher *batchstats.Batcher`
	`147`	`+`
	`148`	`+ WorkspaceAppsStatsCollector *workspaceapps.StatsCollector`
`147`	`149`	`}`
`148`	`150`
`149`	`151`	`// New constructs a codersdk client connected to an in-memory API instance.`
`@@ -425,6 +427,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can`
`425`	`427`	`HealthcheckTimeout: options.HealthcheckTimeout,`
`426`	`428`	`HealthcheckRefresh: options.HealthcheckRefresh,`
`427`	`429`	`StatsBatcher: options.StatsBatcher,`
	`430`	`+ WorkspaceAppsStatsCollector: options.WorkspaceAppsStatsCollector,`
`428`	`431`	`}`
`429`	`432`	`}`
`430`	`433`