From f3b3224cc9b359d987b5c3de0ff1fa133980bd19 Mon Sep 17 00:00:00 2001 From: Gabriel Corado Date: Mon, 6 Jan 2025 12:37:05 -0300 Subject: [PATCH] refactor(test): update option to WithUserAgent --- lib/srv/db/audit_test.go | 2 +- lib/srv/db/common/test.go | 11 +++++------ lib/srv/db/postgres/test.go | 6 +++--- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/lib/srv/db/audit_test.go b/lib/srv/db/audit_test.go index 52a13beb286e2..37aa736fad253 100644 --- 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 index d3e3b725d8784..4eb03e85ae7e1 100644 --- 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 index eb7733b0775b0..b94c7b792d275 100644 --- 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)