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