Skip to content

Commit 7058797

Browse files
kidkidkidkidkidkid
committed
feat(backend): add ut
1 parent 0ecab68 commit 7058797

File tree

3 files changed

+646
-1
lines changed

3 files changed

+646
-1
lines changed

‎backend/modules/observability/infra/repo/ck/annotation_test.go‎

Lines changed: 374 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,374 @@
1+
// Copyright (c) 2025 coze-dev Authors
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package ck
5+
6+
import (
7+
"context"
8+
"errors"
9+
"fmt"
10+
"testing"
11+
12+
"github.com/DATA-DOG/go-sqlmock"
13+
"github.com/coze-dev/coze-loop/backend/modules/observability/infra/repo/ck/gorm_gen/model"
14+
"github.com/stretchr/testify/assert"
15+
"github.com/stretchr/testify/require"
16+
"gorm.io/driver/clickhouse"
17+
"gorm.io/gorm"
18+
)
19+
20+
type mockCkProvider struct {
21+
db *gorm.DB
22+
}
23+
24+
func (m *mockCkProvider) NewSession(ctx context.Context) *gorm.DB {
25+
if m.db == nil {
26+
return nil
27+
}
28+
return m.db.WithContext(ctx)
29+
}
30+
31+
func newInsertAnnotationDao(t *testing.T, failUntil int) (*AnnotationCkDaoImpl, func(), *int) {
32+
t.Helper()
33+
sqlDB, _, err := sqlmock.New()
34+
require.NoError(t, err)
35+
36+
db, err := gorm.Open(clickhouse.New(clickhouse.Config{
37+
Conn: sqlDB,
38+
SkipInitializeWithVersion: true,
39+
}), &gorm.Config{SkipDefaultTransaction: true})
40+
require.NoError(t, err)
41+
42+
count := 0
43+
_ = db.Callback().Create().Replace("gorm:create", func(tx *gorm.DB) {
44+
count++
45+
if count <= failUntil {
46+
tx.Error = errors.New("insert error")
47+
return
48+
}
49+
})
50+
51+
provider := &mockCkProvider{db: db.Session(&gorm.Session{DryRun: true})}
52+
return &AnnotationCkDaoImpl{db: provider}, func() {
53+
_ = sqlDB.Close()
54+
}, &count
55+
}
56+
57+
func newAnnotationDao(t *testing.T) (*AnnotationCkDaoImpl, sqlmock.Sqlmock, *gorm.DB, func()) {
58+
t.Helper()
59+
sqlDB, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
60+
require.NoError(t, err)
61+
62+
db, err := gorm.Open(clickhouse.New(clickhouse.Config{
63+
Conn: sqlDB,
64+
SkipInitializeWithVersion: true,
65+
}), &gorm.Config{SkipDefaultTransaction: true})
66+
require.NoError(t, err)
67+
68+
provider := &mockCkProvider{db: db}
69+
return &AnnotationCkDaoImpl{db: provider}, mock, db, func() {
70+
_ = sqlDB.Close()
71+
}
72+
}
73+
74+
func TestAnnotationCkDaoImpl_Insert(t *testing.T) {
75+
t.Parallel()
76+
77+
ctx := context.Background()
78+
79+
t.Run("nil params", func(t *testing.T) {
80+
dao := &AnnotationCkDaoImpl{}
81+
err := dao.Insert(ctx, nil)
82+
assert.Error(t, err)
83+
})
84+
85+
t.Run("empty annotations", func(t *testing.T) {
86+
dao := &AnnotationCkDaoImpl{}
87+
err := dao.Insert(ctx, &InsertAnnotationParam{})
88+
assert.Error(t, err)
89+
})
90+
91+
t.Run("success", func(t *testing.T) {
92+
dao, cleanup, calls := newInsertAnnotationDao(t, 0)
93+
defer cleanup()
94+
annotation := &model.ObservabilityAnnotation{ID: "anno-1"}
95+
96+
err := dao.Insert(ctx, &InsertAnnotationParam{
97+
Table: "observability_annotations",
98+
Annotations: []*model.ObservabilityAnnotation{annotation},
99+
})
100+
assert.NoError(t, err)
101+
assert.Equal(t, 1, *calls)
102+
})
103+
104+
t.Run("retry failed", func(t *testing.T) {
105+
dao, cleanup, calls := newInsertAnnotationDao(t, 3)
106+
defer cleanup()
107+
annotation := &model.ObservabilityAnnotation{ID: "anno-2"}
108+
109+
err := dao.Insert(ctx, &InsertAnnotationParam{
110+
Table: "observability_annotations",
111+
Annotations: []*model.ObservabilityAnnotation{annotation},
112+
})
113+
assert.Error(t, err)
114+
assert.Equal(t, 3, *calls)
115+
})
116+
}
117+
118+
func TestAnnotationCkDaoImpl_Get(t *testing.T) {
119+
t.Parallel()
120+
121+
ctx := context.Background()
122+
123+
t.Run("invalid params", func(t *testing.T) {
124+
dao := &AnnotationCkDaoImpl{}
125+
_, err := dao.Get(ctx, &GetAnnotationParam{})
126+
assert.Error(t, err)
127+
})
128+
129+
t.Run("build sql error", func(t *testing.T) {
130+
dao, _, _, cleanup := newAnnotationDao(t)
131+
defer cleanup()
132+
133+
_, err := dao.Get(ctx, &GetAnnotationParam{
134+
ID: "anno-1",
135+
StartTime: 1,
136+
EndTime: 2,
137+
})
138+
assert.Error(t, err)
139+
})
140+
141+
t.Run("success", func(t *testing.T) {
142+
dao, mock, _, cleanup := newAnnotationDao(t)
143+
defer cleanup()
144+
145+
rows := sqlmock.NewRows([]string{"id", "span_id"}).AddRow("anno-1", "span-1")
146+
mock.ExpectQuery("SELECT").WillReturnRows(rows)
147+
148+
anno, err := dao.Get(ctx, &GetAnnotationParam{
149+
ID: "anno-1",
150+
Tables: []string{"observability_annotations"},
151+
StartTime: 1,
152+
EndTime: 2,
153+
})
154+
assert.NoError(t, err)
155+
assert.NotNil(t, anno)
156+
assert.Equal(t, "anno-1", anno.ID)
157+
assert.Equal(t, "span-1", anno.SpanID)
158+
assert.NoError(t, mock.ExpectationsWereMet())
159+
})
160+
161+
t.Run("multiple results returns first", func(t *testing.T) {
162+
dao, mock, _, cleanup := newAnnotationDao(t)
163+
defer cleanup()
164+
165+
rows := sqlmock.NewRows([]string{"id", "span_id"}).
166+
AddRow("anno-1", "span-1").
167+
AddRow("anno-2", "span-2")
168+
mock.ExpectQuery("SELECT").WillReturnRows(rows)
169+
170+
anno, err := dao.Get(ctx, &GetAnnotationParam{
171+
ID: "anno-1",
172+
Tables: []string{"observability_annotations"},
173+
StartTime: 1,
174+
EndTime: 2,
175+
})
176+
assert.NoError(t, err)
177+
assert.Equal(t, "anno-1", anno.ID)
178+
assert.NoError(t, mock.ExpectationsWereMet())
179+
})
180+
181+
t.Run("database error", func(t *testing.T) {
182+
dao, mock, _, cleanup := newAnnotationDao(t)
183+
defer cleanup()
184+
185+
mock.ExpectQuery("SELECT").WillReturnError(fmt.Errorf("db error"))
186+
187+
_, err := dao.Get(ctx, &GetAnnotationParam{
188+
ID: "anno-1",
189+
Tables: []string{"observability_annotations"},
190+
StartTime: 1,
191+
EndTime: 2,
192+
})
193+
assert.Error(t, err)
194+
assert.NoError(t, mock.ExpectationsWereMet())
195+
})
196+
}
197+
198+
func TestAnnotationCkDaoImpl_List(t *testing.T) {
199+
t.Parallel()
200+
201+
ctx := context.Background()
202+
203+
t.Run("nil params", func(t *testing.T) {
204+
dao := &AnnotationCkDaoImpl{}
205+
annos, err := dao.List(ctx, nil)
206+
assert.NoError(t, err)
207+
assert.Nil(t, annos)
208+
})
209+
210+
t.Run("empty span ids", func(t *testing.T) {
211+
dao := &AnnotationCkDaoImpl{}
212+
annos, err := dao.List(ctx, &ListAnnotationsParam{})
213+
assert.NoError(t, err)
214+
assert.Nil(t, annos)
215+
})
216+
217+
t.Run("success", func(t *testing.T) {
218+
dao, mock, _, cleanup := newAnnotationDao(t)
219+
defer cleanup()
220+
221+
rows := sqlmock.NewRows([]string{"id", "span_id"}).AddRow("anno-1", "span-1")
222+
mock.ExpectQuery("SELECT").WillReturnRows(rows)
223+
224+
annos, err := dao.List(ctx, &ListAnnotationsParam{
225+
Tables: []string{"observability_annotations"},
226+
SpanIDs: []string{"span-1"},
227+
StartTime: 1,
228+
EndTime: 2,
229+
Limit: 10,
230+
})
231+
assert.NoError(t, err)
232+
require.Len(t, annos, 1)
233+
assert.Equal(t, "span-1", annos[0].SpanID)
234+
assert.NoError(t, mock.ExpectationsWereMet())
235+
})
236+
237+
t.Run("database error", func(t *testing.T) {
238+
dao, mock, _, cleanup := newAnnotationDao(t)
239+
defer cleanup()
240+
241+
mock.ExpectQuery("SELECT").WillReturnError(fmt.Errorf("db error"))
242+
243+
_, err := dao.List(ctx, &ListAnnotationsParam{
244+
Tables: []string{"observability_annotations"},
245+
SpanIDs: []string{"span-1"},
246+
StartTime: 1,
247+
EndTime: 2,
248+
Limit: 10,
249+
})
250+
assert.Error(t, err)
251+
assert.NoError(t, mock.ExpectationsWereMet())
252+
})
253+
254+
t.Run("build sql error", func(t *testing.T) {
255+
dao, _, _, cleanup := newAnnotationDao(t)
256+
defer cleanup()
257+
258+
_, err := dao.List(ctx, &ListAnnotationsParam{
259+
SpanIDs: []string{"span-1"},
260+
StartTime: 1,
261+
EndTime: 2,
262+
})
263+
assert.Error(t, err)
264+
})
265+
}
266+
267+
func TestAnnotationCkDaoImpl_buildSingleSql(t *testing.T) {
268+
t.Parallel()
269+
270+
dao, _, db, cleanup := newAnnotationDao(t)
271+
defer cleanup()
272+
273+
ctx := context.Background()
274+
baseSession := dao.db.NewSession(ctx)
275+
require.NotNil(t, baseSession)
276+
277+
testCases := []struct {
278+
name string
279+
param *annoSqlParam
280+
assert func(t *testing.T, sql string)
281+
}{
282+
{
283+
name: "with id filter",
284+
param: &annoSqlParam{
285+
Tables: []string{"observability_annotations"},
286+
ID: "anno-1",
287+
StartTime: 1,
288+
EndTime: 2,
289+
Limit: 10,
290+
},
291+
assert: func(t *testing.T, sql string) {
292+
assert.Contains(t, sql, "FROM `observability_annotations`")
293+
assert.Contains(t, sql, "id = 'anno-1'")
294+
assert.Contains(t, sql, "ORDER BY `created_at`")
295+
assert.Contains(t, sql, "LIMIT 10")
296+
},
297+
},
298+
{
299+
name: "with span ids and desc",
300+
param: &annoSqlParam{
301+
Tables: []string{"observability_annotations"},
302+
SpanIDs: []string{"span-1"},
303+
StartTime: 10,
304+
EndTime: 20,
305+
Limit: 5,
306+
DescByUpdatedAt: true,
307+
},
308+
assert: func(t *testing.T, sql string) {
309+
assert.Contains(t, sql, "span_id IN ('span-1')")
310+
assert.Contains(t, sql, "ORDER BY `updated_at` DESC")
311+
assert.Contains(t, sql, "LIMIT 5")
312+
},
313+
},
314+
}
315+
316+
for _, tc := range testCases {
317+
tc := tc
318+
t.Run(tc.name, func(t *testing.T) {
319+
session := baseSession.Session(&gorm.Session{DryRun: true})
320+
query, err := dao.buildSingleSql(ctx, session, tc.param.Tables[0], tc.param)
321+
require.NoError(t, err)
322+
sql := query.ToSQL(func(tx *gorm.DB) *gorm.DB {
323+
return tx.Find([]*model.ObservabilityAnnotation{})
324+
})
325+
tc.assert(t, sql)
326+
})
327+
}
328+
329+
_ = db // silence unused (db kept for completeness)
330+
}
331+
332+
func TestAnnotationCkDaoImpl_buildSql(t *testing.T) {
333+
t.Parallel()
334+
335+
dao, _, _, cleanup := newAnnotationDao(t)
336+
defer cleanup()
337+
338+
ctx := context.Background()
339+
340+
t.Run("no tables", func(t *testing.T) {
341+
_, err := dao.buildSql(ctx, &annoSqlParam{})
342+
assert.Error(t, err)
343+
})
344+
345+
t.Run("single table", func(t *testing.T) {
346+
result, err := dao.buildSql(ctx, &annoSqlParam{
347+
Tables: []string{"observability_annotations"},
348+
StartTime: 1,
349+
EndTime: 2,
350+
Limit: 3,
351+
})
352+
assert.NoError(t, err)
353+
sql := result.Statement.SQL.String()
354+
assert.Contains(t, sql, "FROM `observability_annotations`")
355+
assert.Contains(t, sql, "LIMIT 3")
356+
assert.Contains(t, sql, "SETTINGS final = 1")
357+
})
358+
359+
t.Run("multiple tables", func(t *testing.T) {
360+
result, err := dao.buildSql(ctx, &annoSqlParam{
361+
Tables: []string{"observability_annotations", "observability_annotations_v2"},
362+
StartTime: 1,
363+
EndTime: 2,
364+
Limit: 5,
365+
DescByUpdatedAt: true,
366+
})
367+
assert.NoError(t, err)
368+
sql := result.Statement.SQL.String()
369+
assert.Contains(t, sql, "UNION ALL")
370+
assert.Contains(t, sql, "ORDER BY updated_at DESC")
371+
assert.Contains(t, sql, "LIMIT 5")
372+
assert.Contains(t, sql, "SETTINGS final = 1")
373+
})
374+
}

‎backend/modules/observability/infra/repo/ck/spans.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,9 @@ func (s *SpansCkDaoImpl) Get(ctx context.Context, param *QueryParam) ([]*model.O
115115
return nil, errorx.WrapByCode(err, obErrorx.CommercialCommonRPCErrorCodeCode)
116116
}
117117
for _, span := range spans {
118+
if span.SystemTagsString == nil {
119+
span.SystemTagsString = make(map[string]string)
120+
}
118121
span.SystemTagsString[loop_span.SpanFieldTenant] = "cozeloop" // tenant
119122
}
120123
return spans, nil

0 commit comments

Comments
 (0)