1 package sources_test
2
3 import (
4 "context"
5 "testing"
6
7 "github.com/jackc/pgx/v5/pgconn"
8 "github.com/jackc/pgx/v5/pgxpool"
9 "github.com/pashagolub/pgxmock/v4"
10 "github.com/stretchr/testify/assert"
11 "github.com/stretchr/testify/require"
12
13 "github.com/cybertec-postgresql/pgwatch/v3/internal/db"
14 "github.com/cybertec-postgresql/pgwatch/v3/internal/sources"
15 )
16
17 const ImageName = "docker.io/postgres:17-alpine"
18
19 func TestSourceConn_Connect(t *testing.T) {
20
21 t.Run("failed config parsing", func(t *testing.T) {
22 md := &sources.SourceConn{}
23 md.ConnStr = "invalid connection string"
24 err := md.Connect(ctx, sources.CmdOpts{})
25 assert.Error(t, err)
26 })
27
28 t.Run("failed connection", func(t *testing.T) {
29 md := &sources.SourceConn{}
30 sources.NewConnWithConfig = func(_ context.Context, _ *pgxpool.Config, _ ...db.ConnConfigCallback) (db.PgxPoolIface, error) {
31 return nil, assert.AnError
32 }
33 err := md.Connect(ctx, sources.CmdOpts{})
34 assert.ErrorIs(t, err, assert.AnError)
35 })
36
37 t.Run("successful connection to pgbouncer", func(t *testing.T) {
38 mock, err := pgxmock.NewPool()
39 require.NoError(t, err)
40 sources.NewConnWithConfig = func(_ context.Context, _ *pgxpool.Config, _ ...db.ConnConfigCallback) (db.PgxPoolIface, error) {
41 return mock, nil
42 }
43
44 md := &sources.SourceConn{}
45 md.Kind = sources.SourcePgBouncer
46
47 opts := sources.CmdOpts{}
48 opts.MaxParallelConnectionsPerDb = 3
49
50 mock.ExpectExec("SHOW VERSION").WillReturnResult(pgconn.NewCommandTag("SELECT 1"))
51
52 err = md.Connect(ctx, opts)
53 assert.NoError(t, err)
54
55 assert.NoError(t, mock.ExpectationsWereMet())
56 })
57 }
58
59 func TestSourceConn_ParseConfig(t *testing.T) {
60 md := &sources.SourceConn{}
61 assert.NoError(t, md.ParseConfig())
62
63 assert.NoError(t, md.ParseConfig())
64 }
65
66 func TestSourceConn_GetDatabaseName(t *testing.T) {
67 md := &sources.SourceConn{}
68 md.ConnStr = "postgres://user:password@localhost:5432/mydatabase"
69 expected := "mydatabase"
70
71 got := md.GetDatabaseName()
72 assert.Equal(t, expected, got, "GetDatabaseName() = %v, want %v", got, expected)
73
74 got = md.Source.GetDatabaseName()
75 assert.Equal(t, expected, got, "GetDatabaseName() = %v, want %v", got, expected)
76
77 md = &sources.SourceConn{}
78 md.ConnStr = "foo boo"
79 expected = ""
80 got = md.GetDatabaseName()
81 assert.Equal(t, expected, got, "GetDatabaseName() = %v, want %v", got, expected)
82 }
83
84 func TestSourceConn_SetDatabaseName(t *testing.T) {
85 md := &sources.SourceConn{}
86 md.ConnStr = "postgres://user:password@localhost:5432/mydatabase"
87 expected := "mydatabase"
88
89 md.SetDatabaseName(expected)
90 got := md.GetDatabaseName()
91 assert.Equal(t, expected, got, "GetDatabaseName() = %v, want %v", got, expected)
92
93 expected = "newdatabase"
94 md.SetDatabaseName(expected)
95 got = md.GetDatabaseName()
96 assert.Equal(t, expected, got, "GetDatabaseName() = %v, want %v", got, expected)
97
98 md = &sources.SourceConn{}
99 md.ConnStr = "foo boo"
100 expected = ""
101 md.SetDatabaseName("ingored due to invalid ConnStr")
102 got = md.GetDatabaseName()
103 assert.Equal(t, expected, got, "GetDatabaseName() = %v, want %v", got, expected)
104 }
105
106 func TestSourceConn_DiscoverPlatform(t *testing.T) {
107 mock, err := pgxmock.NewPool()
108 require.NoError(t, err)
109 md := &sources.SourceConn{Conn: mock}
110
111 mock.ExpectQuery("select").WillReturnRows(pgxmock.NewRows([]string{"exec_env"}).AddRow("AZURE_SINGLE"))
112
113 assert.Equal(t, "AZURE_SINGLE", md.DiscoverPlatform(ctx))
114 assert.NoError(t, mock.ExpectationsWereMet())
115 }
116
117 func TestSourceConn_GetApproxSize(t *testing.T) {
118 mock, err := pgxmock.NewPool()
119 require.NoError(t, err)
120 md := &sources.SourceConn{Conn: mock}
121
122 mock.ExpectQuery("select").WillReturnRows(pgxmock.NewRows([]string{"size"}).AddRow(42))
123
124 size, err := md.GetApproxSize(ctx)
125 assert.EqualValues(t, 42, size)
126 assert.NoError(t, err)
127 assert.NoError(t, mock.ExpectationsWereMet())
128 }
129
130 func TestSourceConn_FunctionExists(t *testing.T) {
131 mock, err := pgxmock.NewPool()
132 require.NoError(t, err)
133 md := &sources.SourceConn{Conn: mock}
134
135 mock.ExpectQuery("select").WithArgs("get_foo").WillReturnRows(pgxmock.NewRows([]string{"exists"}))
136
137 assert.False(t, md.FunctionExists(ctx, "get_foo"))
138 assert.NoError(t, mock.ExpectationsWereMet())
139 }
140
141 func TestSourceConn_IsPostgresSource(t *testing.T) {
142 md := &sources.SourceConn{}
143 md.Kind = sources.SourcePostgres
144 assert.True(t, md.IsPostgresSource(), "IsPostgresSource() = false, want true")
145
146 md.Kind = sources.SourcePgBouncer
147 assert.False(t, md.IsPostgresSource(), "IsPostgresSource() = true, want false")
148
149 md.Kind = sources.SourcePgPool
150 assert.False(t, md.IsPostgresSource(), "IsPostgresSource() = true, want false")
151
152 md.Kind = sources.SourcePatroni
153 assert.True(t, md.IsPostgresSource(), "IsPostgresSource() = false, want true")
154 }
155
156 func TestSourceConn_Ping(t *testing.T) {
157 db, err := pgxmock.NewPool()
158 require.NoError(t, err)
159 md := &sources.SourceConn{Conn: db}
160
161 db.ExpectPing()
162 md.Kind = sources.SourcePostgres
163 assert.NoError(t, md.Ping(ctx), "Ping() = error, want nil")
164
165 db.ExpectExec("SHOW VERSION").WillReturnResult(pgconn.NewCommandTag("SELECT 1"))
166 md.Conn = db
167 md.Kind = sources.SourcePgBouncer
168 assert.NoError(t, md.Ping(ctx), "Ping() = error, want nil")
169 }
170
171 type testSourceReader struct {
172 sources.Sources
173 error
174 }
175
176 func (r testSourceReader) GetSources() (sources.Sources, error) {
177 return r.Sources, r.error
178 }
179
180 func TestMonitoredDatabases_SyncFromReader_error(t *testing.T) {
181 reader := testSourceReader{error: assert.AnError}
182 mds := sources.SourceConns{}
183 _, err := mds.SyncFromReader(reader)
184 assert.Error(t, err)
185 }
186
187 func TestMonitoredDatabases_SyncFromReader(t *testing.T) {
188 db, _ := pgxmock.NewPool()
189 src := sources.Source{
190 Name: "test",
191 Kind: sources.SourcePostgres,
192 IsEnabled: true,
193 ConnStr: "postgres://user:password@localhost:5432/mydatabase",
194 }
195 reader := testSourceReader{Sources: sources.Sources{src}}
196
197 mds, _ := reader.GetSources()
198 assert.NotNil(t, mds, "GetSources() = nil, want not nil")
199
200 mdbs, _ := mds.ResolveDatabases()
201 assert.NotNil(t, mdbs, "ResolveDatabases() = nil, want not nil")
202
203 mdbs[0].Conn = db
204 db.ExpectClose()
205
206 newmdbs, _ := mdbs.SyncFromReader(reader)
207 assert.NotNil(t, newmdbs)
208 assert.Equal(t, mdbs[0].ConnStr, newmdbs[0].ConnStr)
209 assert.Equal(t, db, newmdbs[0].Conn)
210
211 reader.Sources[0].ConnStr = "postgres://user:password@localhost:5432/anotherdatabase"
212 newmdbs, _ = mdbs.SyncFromReader(reader)
213 assert.NotNil(t, newmdbs)
214 assert.NotEqual(t, mdbs[0].ConnStr, newmdbs[0].ConnStr)
215 assert.Nil(t, newmdbs[0].Conn)
216 assert.NoError(t, db.ExpectationsWereMet())
217
218 reader.Sources[0].Name = "another"
219 newmdbs, _ = mdbs.SyncFromReader(reader)
220 assert.NotNil(t, newmdbs)
221 assert.NotEqual(t, mdbs[0].Name, newmdbs[0].Name)
222 assert.Nil(t, newmdbs[0].Conn)
223 }
224