refactor(test): update option to WithUserAgent

gravitational · Jan 6, 2025 · f3b3224 · f3b3224
1 parent aaae0ed
commit f3b3224
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 10 deletions.
diff --git a/lib/srv/db/audit_test.go b/lib/srv/db/audit_test.go
@@ -61,7 +61,7 @@ func TestAuditPostgres(t *testing.T) {
 
 	// Connect should trigger successful session start event.
 	userAgent := "psql"
-	psql, err := testCtx.postgresClient(ctx, "alice", "postgres", "postgres", "postgres", common.WithParams(map[string]string{"application_name": userAgent}))
+	psql, err := testCtx.postgresClient(ctx, "alice", "postgres", "postgres", "postgres", common.WithUserAgent(userAgent))
 	require.NoError(t, err)
 	startEvt, ok := requireEvent(t, testCtx, libevents.DatabaseSessionStartCode).(*events.DatabaseSessionStart)
 	require.True(t, ok)

diff --git a/lib/srv/db/common/test.go b/lib/srv/db/common/test.go
@@ -163,10 +163,10 @@ func MakeTestServerTLSConfig(config TestServerConfig) (*tls.Config, error) {
 // ClientOption represents a database client config option.
 type ClientOption func(config *TestClientConfig)
 
-// WithParams set client config params.
-func WithParams(params map[string]string) ClientOption {
+// WithUserAgent set client user agent.
+func WithUserAgent(userAgent string) ClientOption {
 	return func(config *TestClientConfig) {
-		config.Params = params
+		config.UserAgent = userAgent
 	}
 }
 
@@ -186,9 +186,8 @@ type TestClientConfig struct {
 	PinnedIP string
 	// RouteToDatabase contains database routing information.
 	RouteToDatabase tlsca.RouteToDatabase
-	// Params contains parameters used during the session. This can hold, for
-	// example, PostgreSQL runtime parameters.
-	Params map[string]string
+	// UserAgent contains the client user agent.
+	UserAgent string
 }
 
 // MakeTestClientTLSCert returns TLS certificate suitable for configuring test

diff --git a/lib/srv/db/postgres/test.go b/lib/srv/db/postgres/test.go
@@ -55,12 +55,12 @@ func MakeTestClient(ctx context.Context, config common.TestClientConfig) (*pgcon
 	pgconnConfig.User = config.RouteToDatabase.Username
 	pgconnConfig.Database = config.RouteToDatabase.Database
 	pgconnConfig.TLSConfig, err = common.MakeTestClientTLSConfig(config)
-	for name, value := range config.Params {
-		pgconnConfig.RuntimeParams[name] = value
-	}
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
+	if config.UserAgent != "" {
+		pgconnConfig.RuntimeParams["application_name"] = config.UserAgent
+	}
 	pgConn, err := pgconn.ConnectConfig(ctx, pgconnConfig)
 	if err != nil {
 		return nil, trace.Wrap(err)