1 package sinks_test
2
3 import (
4 "testing"
5
6 "github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
7 "github.com/cybertec-postgresql/pgwatch/v6/internal/sinks"
8 "github.com/cybertec-postgresql/pgwatch/v6/internal/testutil"
9 "github.com/stretchr/testify/assert"
10 )
11
12
13 type mockWriter struct {
14 err error
15 needsMigration bool
16 needsMigrationErr error
17 }
18
19 func (m *mockWriter) SyncMetric(string, string, sinks.SyncOp) error {
20 return m.err
21 }
22
23 func (m *mockWriter) Write(metrics.MeasurementEnvelope) error {
24 return m.err
25 }
26
27 func (m *mockWriter) Migrate() error {
28 return m.err
29 }
30
31 func (m *mockWriter) NeedsMigration() (bool, error) {
32 return m.needsMigration, m.needsMigrationErr
33 }
34
35 func (m *mockWriter) DefineMetrics(*metrics.Metrics) error {
36 return m.err
37 }
38
39 func TestNewMultiWriter(t *testing.T) {
40 input := []struct {
41 opts *sinks.CmdOpts
42 w bool
43 err bool
44 }{
45 {&sinks.CmdOpts{}, false, true},
46 {&sinks.CmdOpts{
47 Sinks: []string{"foo"},
48 }, false, true},
49 {&sinks.CmdOpts{
50 Sinks: []string{"jsonfile://test.json"},
51 }, true, false},
52 {&sinks.CmdOpts{
53 Sinks: []string{"jsonfile://test.json", "jsonfile://test1.json"},
54 }, true, false},
55 {&sinks.CmdOpts{
56 Sinks: []string{"prometheus://foo/"},
57 }, false, true},
58 {&sinks.CmdOpts{
59 Sinks: []string{"rpc://foo/"},
60 }, false, true},
61 {&sinks.CmdOpts{
62 Sinks: []string{"postgresql:///baz"},
63 }, false, true},
64 {&sinks.CmdOpts{
65 Sinks: []string{"foo:///"},
66 }, false, true},
67 }
68
69 for _, i := range input {
70 mw, err := sinks.NewSinkWriter(testutil.TestContext, i.opts)
71 if i.err {
72 assert.Error(t, err)
73 } else {
74 assert.NoError(t, err)
75 }
76 if i.w {
77 assert.NotNil(t, mw)
78 } else {
79 assert.Nil(t, mw)
80 }
81 }
82 }
83
84 func TestAddWriter(t *testing.T) {
85 mw := &sinks.MultiWriter{}
86 mockWriter := &mockWriter{}
87 mw.AddWriter(mockWriter)
88 assert.Equal(t, 1, mw.Count())
89 }
90
91 func TestSyncMetrics(t *testing.T) {
92 mw := &sinks.MultiWriter{}
93 mockWriter := &mockWriter{}
94 mw.AddWriter(mockWriter)
95 err := mw.SyncMetric("db", "metric", sinks.InvalidOp)
96 assert.NoError(t, err)
97 }
98
99 func TestWriteMeasurements(t *testing.T) {
100 mw := &sinks.MultiWriter{}
101 mockWriter := &mockWriter{}
102 mw.AddWriter(mockWriter)
103 err := mw.Write(metrics.MeasurementEnvelope{})
104 assert.NoError(t, err)
105 }
106
107 func TestMultiWriterMigrate(t *testing.T) {
108 tests := []struct {
109 name string
110 writers []sinks.Writer
111 expectError bool
112 }{
113 {
114 name: "no migratable writers",
115 writers: []sinks.Writer{
116 &mockWriter{},
117 },
118 expectError: false,
119 },
120 {
121 name: "single migratable writer success",
122 writers: []sinks.Writer{
123 &mockWriter{},
124 },
125 expectError: false,
126 },
127 {
128 name: "single migratable writer error",
129 writers: []sinks.Writer{
130 &mockWriter{err: assert.AnError},
131 },
132 expectError: true,
133 },
134 {
135 name: "multiple migratable writers success",
136 writers: []sinks.Writer{
137 &mockWriter{},
138 &mockWriter{},
139 },
140 expectError: false,
141 },
142 {
143 name: "multiple writers with one error",
144 writers: []sinks.Writer{
145 &mockWriter{},
146 &mockWriter{err: assert.AnError},
147 },
148 expectError: true,
149 },
150 {
151 name: "mixed writers with migration error",
152 writers: []sinks.Writer{
153 &mockWriter{},
154 &mockWriter{err: assert.AnError},
155 },
156 expectError: true,
157 },
158 }
159
160 for _, tt := range tests {
161 t.Run(tt.name, func(t *testing.T) {
162 mw := &sinks.MultiWriter{}
163 for _, w := range tt.writers {
164 mw.AddWriter(w)
165 }
166 err := mw.Migrate()
167 if tt.expectError {
168 assert.Error(t, err)
169 } else {
170 assert.NoError(t, err)
171 }
172 })
173 }
174 }
175
176 func TestDefineMetrics(t *testing.T) {
177 tests := []struct {
178 name string
179 writers []sinks.Writer
180 expectError bool
181 }{
182 {
183 name: "writer without DefineMetrics",
184 writers: []sinks.Writer{&mockWriter{}},
185 expectError: false,
186 },
187 {
188 name: "single definer success",
189 writers: []sinks.Writer{&mockWriter{err: nil}},
190 expectError: false,
191 },
192 {
193 name: "single definer error",
194 writers: []sinks.Writer{&mockWriter{err: assert.AnError}},
195 expectError: true,
196 },
197 {
198 name: "two definers errors",
199 writers: []sinks.Writer{
200 &mockWriter{err: assert.AnError},
201 &mockWriter{err: assert.AnError},
202 },
203 expectError: true,
204 },
205 {
206 name: "mixed writers error",
207 writers: []sinks.Writer{
208 &mockWriter{},
209 &mockWriter{err: assert.AnError},
210 },
211 expectError: true,
212 },
213 {
214 name: "mixed writers success",
215 writers: []sinks.Writer{
216 &mockWriter{},
217 &mockWriter{err: nil},
218 },
219 expectError: false,
220 },
221 }
222
223 for _, tt := range tests {
224 t.Run(tt.name, func(t *testing.T) {
225 mw := &sinks.MultiWriter{}
226 for _, w := range tt.writers {
227 mw.AddWriter(w)
228 }
229 err := mw.DefineMetrics(&metrics.Metrics{})
230 if tt.expectError {
231 assert.Error(t, err)
232 } else {
233 assert.NoError(t, err)
234 }
235 })
236 }
237 }
238
239 func TestMultiWriterNeedsMigration(t *testing.T) {
240 tests := []struct {
241 name string
242 writers []sinks.Writer
243 expectNeedsMigrate bool
244 expectError bool
245 }{
246 {
247 name: "no migratable writers",
248 writers: []sinks.Writer{
249 &mockWriter{},
250 },
251 expectNeedsMigrate: false,
252 expectError: false,
253 },
254 {
255 name: "single writer needs migration",
256 writers: []sinks.Writer{
257 &mockWriter{needsMigration: true},
258 },
259 expectNeedsMigrate: true,
260 expectError: false,
261 },
262 {
263 name: "single writer no migration needed",
264 writers: []sinks.Writer{
265 &mockWriter{needsMigration: false},
266 },
267 expectNeedsMigrate: false,
268 expectError: false,
269 },
270 {
271 name: "multiple writers one needs migration",
272 writers: []sinks.Writer{
273 &mockWriter{needsMigration: false},
274 &mockWriter{needsMigration: true},
275 },
276 expectNeedsMigrate: true,
277 expectError: false,
278 },
279 {
280 name: "error checking migration",
281 writers: []sinks.Writer{
282 &mockWriter{needsMigrationErr: assert.AnError},
283 },
284 expectNeedsMigrate: false,
285 expectError: true,
286 },
287 {
288 name: "mixed writers one needs migration",
289 writers: []sinks.Writer{
290 &mockWriter{},
291 &mockWriter{needsMigration: true},
292 },
293 expectNeedsMigrate: true,
294 expectError: false,
295 },
296 }
297
298 for _, tt := range tests {
299 t.Run(tt.name, func(t *testing.T) {
300 mw := &sinks.MultiWriter{}
301 for _, w := range tt.writers {
302 mw.AddWriter(w)
303 }
304 needs, err := mw.NeedsMigration()
305 assert.Equal(t, tt.expectNeedsMigrate, needs)
306 if tt.expectError {
307 assert.Error(t, err)
308 } else {
309 assert.NoError(t, err)
310 }
311 })
312 }
313 }
314