1 package metrics
2
3 import (
4 "context"
5 "testing"
6
7 migrator "github.com/cybertec-postgresql/pgx-migrator"
8 "github.com/pashagolub/pgxmock/v5"
9 "github.com/stretchr/testify/assert"
10 )
11
12 func TestMigrate(t *testing.T) {
13 a := assert.New(t)
14 conn, err := pgxmock.NewPool()
15 a.NoError(err)
16
17 conn.ExpectExec(`CREATE TABLE IF NOT EXISTS pgwatch\.migration`).WillReturnResult(pgxmock.NewResult("CREATE", 1))
18 conn.ExpectQuery(`SELECT count`).WillReturnRows(pgxmock.NewRows([]string{"count"}).AddRow(0))
19 conn.ExpectBegin()
20 conn.ExpectExec(`INSERT INTO`).WillReturnResult(pgxmock.NewResult("INSERT", 1))
21 conn.ExpectBegin()
22 conn.ExpectExec(`UPDATE pgwatch\.metric`).WillReturnResult(pgxmock.NewResult("UPDATE", 0))
23 conn.ExpectExec(`INSERT INTO`).WillReturnResult(pgxmock.NewResult("INSERT", 1))
24 conn.ExpectBegin()
25 conn.ExpectExec(`ALTER TABLE`).WillReturnResult(pgxmock.NewResult("ALTER TABLE", 0))
26 conn.ExpectExec(`INSERT INTO`).WillReturnResult(pgxmock.NewResult("INSERT", 1))
27
28 dmrw := &dbMetricReaderWriter{ctx, conn}
29 err = dmrw.Migrate()
30 a.NoError(err)
31 }
32
33 func TestNeedsMigration(t *testing.T) {
34 a := assert.New(t)
35 conn, err := pgxmock.NewPool()
36 a.NoError(err)
37
38 conn.ExpectQuery(`SELECT to_regclass`).
39 WithArgs("pgwatch.migration").
40 WillReturnRows(pgxmock.NewRows([]string{"to_regclass"}).AddRow(true))
41 conn.ExpectQuery(`SELECT count`).WillReturnRows(pgxmock.NewRows([]string{"count"}).AddRow(0))
42
43 dmrw := &dbMetricReaderWriter{ctx, conn}
44 needs, err := dmrw.NeedsMigration()
45 a.NoError(err)
46 a.True(needs)
47 }
48
49 func TestMigrateFail(t *testing.T) {
50 oldInitMigrator := initMigrator
51 t.Cleanup(func() {
52 initMigrator = oldInitMigrator
53 })
54 a := assert.New(t)
55 dmrw := &dbMetricReaderWriter{}
56 initMigrator = func(context.Context) (*migrator.Migrator, error) {
57 return nil, assert.AnError
58 }
59 err := dmrw.Migrate()
60 a.Error(err)
61 }
62
63 func TestNeedsMigrationFail(t *testing.T) {
64 oldInitMigrator := initMigrator
65 t.Cleanup(func() {
66 initMigrator = oldInitMigrator
67 })
68 a := assert.New(t)
69 dmrw := &dbMetricReaderWriter{}
70 initMigrator = func(context.Context) (*migrator.Migrator, error) {
71 return nil, assert.AnError
72 }
73 _, err := dmrw.NeedsMigration()
74 a.Error(err)
75 }
76