Skip to content

Commit ae3495a

Browse files
henrybarretohenrybarreto
committed
feature(api): add migration to delete recorded_sessions collections
1 parent 3377b63 commit ae3495a

File tree

3 files changed

+311
-0
lines changed

3 files changed

+311
-0
lines changed

‎api/store/mongo/migrations/main.go

+1
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ func GenerateMigrations() []migrate.Migration {
104104
migration92,
105105
migration93,
106106
migration94,
107+
migration95,
107108
}
108109
}
109110

‎api/store/mongo/migrations/migration_95.go

+49
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package migrations
2+
3+
import (
4+
"context"
5+
6+
"github.com/shellhub-io/shellhub/pkg/envs"
7+
log "github.com/sirupsen/logrus"
8+
migrate "github.com/xakep666/mongo-migrate"
9+
"go.mongodb.org/mongo-driver/mongo"
10+
)
11+
12+
var migration95 = migrate.Migration{
13+
Version: 95,
14+
Description: "Drops the recorded_session collection",
15+
Up: migrate.MigrationFunc(func(ctx context.Context, db *mongo.Database) error {
16+
log.WithFields(log.Fields{
17+
"component": "migration",
18+
"version": 95,
19+
"action": "Up",
20+
}).Info("Applying migration")
21+
22+
if !envs.IsEnterprise() {
23+
return nil
24+
}
25+
26+
if err := db.Collection("recorded_sessions").Drop(ctx); err != nil {
27+
return err
28+
}
29+
30+
return nil
31+
}),
32+
Down: migrate.MigrationFunc(func(ctx context.Context, db *mongo.Database) error {
33+
log.WithFields(log.Fields{
34+
"component": "migration",
35+
"version": 95,
36+
"action": "Down",
37+
}).Info("Reverting migration")
38+
39+
if !envs.IsEnterprise() {
40+
return nil
41+
}
42+
43+
if err := db.CreateCollection(ctx, "recorded_sessions"); err != nil {
44+
return err
45+
}
46+
47+
return nil
48+
}),
49+
}

‎api/store/mongo/migrations/migration_95_test.go

+261
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,261 @@
1+
package migrations
2+
3+
import (
4+
"context"
5+
"slices"
6+
"testing"
7+
8+
"github.com/shellhub-io/shellhub/pkg/envs"
9+
envmock "github.com/shellhub-io/shellhub/pkg/envs/mocks"
10+
"github.com/stretchr/testify/assert"
11+
"github.com/stretchr/testify/require"
12+
migrate "github.com/xakep666/mongo-migrate"
13+
"go.mongodb.org/mongo-driver/bson"
14+
)
15+
16+
func TestMigration95Up(t *testing.T) {
17+
ctx := context.Background()
18+
mock := &envmock.Backend{}
19+
envs.DefaultBackend = mock
20+
21+
cases := []struct {
22+
description string
23+
setup func() error
24+
verify func(tt *testing.T)
25+
}{
26+
{
27+
description: "Drop recorded_sessions collection in enterprise mode",
28+
setup: func() error {
29+
mock.On("Get", "SHELLHUB_ENTERPRISE").Return("true").Once()
30+
31+
_, err := c.
32+
Database("test").
33+
Collection("sessions").
34+
InsertOne(ctx, bson.M{
35+
"uid": "session-1",
36+
"authenticated": true,
37+
})
38+
if err != nil {
39+
return err
40+
}
41+
42+
_, err = c.
43+
Database("test").
44+
Collection("recorded_sessions").
45+
InsertMany(ctx, []any{
46+
bson.M{
47+
"uid": "session-1",
48+
"message": "test output 1",
49+
"time": "2023-01-01T10:00:00Z",
50+
"width": 80,
51+
"height": 24,
52+
},
53+
bson.M{
54+
"uid": "session-1",
55+
"message": "test output 2",
56+
"time": "2023-01-01T10:01:00Z",
57+
"width": 80,
58+
"height": 24,
59+
},
60+
})
61+
62+
return err
63+
},
64+
verify: func(tt *testing.T) {
65+
collections, err := c.Database("test").ListCollectionNames(ctx, bson.M{})
66+
require.NoError(tt, err)
67+
68+
hasRecordedSessions := slices.Contains(collections, "recorded_sessions")
69+
70+
if hasRecordedSessions {
71+
count, err := c.
72+
Database("test").
73+
Collection("recorded_sessions").
74+
CountDocuments(ctx, bson.M{})
75+
require.NoError(tt, err)
76+
assert.Equal(tt, int64(0), count)
77+
}
78+
79+
count, err := c.
80+
Database("test").
81+
Collection("sessions").
82+
CountDocuments(ctx, bson.M{})
83+
require.NoError(tt, err)
84+
assert.Equal(tt, int64(1), count)
85+
},
86+
},
87+
{
88+
description: "Nothing happens when not in enterprise mode",
89+
setup: func() error {
90+
mock.On("Get", "SHELLHUB_ENTERPRISE").Return("false").Once()
91+
92+
_, err := c.
93+
Database("test").
94+
Collection("sessions").
95+
InsertOne(ctx, bson.M{
96+
"uid": "session-2",
97+
"authenticated": true,
98+
})
99+
if err != nil {
100+
return err
101+
}
102+
103+
_, err = c.
104+
Database("test").
105+
Collection("recorded_sessions").
106+
InsertOne(ctx, bson.M{
107+
"uid": "session-2",
108+
"message": "test output",
109+
"time": "2023-01-01T10:00:00Z",
110+
"width": 80,
111+
"height": 24,
112+
})
113+
114+
return err
115+
},
116+
verify: func(tt *testing.T) {
117+
count, err := c.
118+
Database("test").
119+
Collection("recorded_sessions").
120+
CountDocuments(ctx, bson.M{"uid": "session-2"})
121+
require.NoError(tt, err)
122+
assert.Equal(tt, int64(1), count)
123+
},
124+
},
125+
}
126+
127+
for _, tc := range cases {
128+
t.Run(tc.description, func(tt *testing.T) {
129+
tt.Cleanup(func() {
130+
assert.NoError(tt, srv.Reset())
131+
})
132+
133+
require.NoError(tt, tc.setup())
134+
135+
migrates := migrate.NewMigrate(c.Database("test"), GenerateMigrations()[94])
136+
require.NoError(tt, migrates.Up(ctx, migrate.AllAvailable))
137+
138+
tc.verify(tt)
139+
})
140+
}
141+
}
142+
143+
func TestMigration95Down(t *testing.T) {
144+
ctx := context.Background()
145+
mock := &envmock.Backend{}
146+
envs.DefaultBackend = mock
147+
148+
cases := []struct {
149+
description string
150+
setup func() error
151+
verify func(tt *testing.T)
152+
}{
153+
{
154+
description: "Migration down has no effect",
155+
setup: func() error {
156+
mock.On("Get", "SHELLHUB_ENTERPRISE").Return("true").Twice()
157+
158+
_, err := c.
159+
Database("test").
160+
Collection("sessions").
161+
InsertOne(ctx, bson.M{
162+
"uid": "session-3",
163+
"authenticated": true,
164+
})
165+
if err != nil {
166+
return err
167+
}
168+
169+
migrates := migrate.NewMigrate(c.Database("test"), GenerateMigrations()[94])
170+
if err := migrates.Up(ctx, migrate.AllAvailable); err != nil {
171+
return err
172+
}
173+
174+
return nil
175+
},
176+
verify: func(tt *testing.T) {
177+
collections, err := c.Database("test").ListCollectionNames(ctx, bson.M{})
178+
require.NoError(tt, err)
179+
180+
hasRecordedSessions := slices.Contains(collections, "recorded_sessions")
181+
182+
if hasRecordedSessions {
183+
count, err := c.
184+
Database("test").
185+
Collection("recorded_sessions").
186+
CountDocuments(ctx, bson.M{})
187+
require.NoError(tt, err)
188+
assert.Equal(tt, int64(0), count)
189+
}
190+
191+
count, err := c.
192+
Database("test").
193+
Collection("sessions").
194+
CountDocuments(ctx, bson.M{})
195+
require.NoError(tt, err)
196+
assert.Equal(tt, int64(1), count)
197+
},
198+
},
199+
{
200+
description: "Nothing happens when not in enterprise mode",
201+
setup: func() error {
202+
mock.On("Get", "SHELLHUB_ENTERPRISE").Return("false").Twice()
203+
204+
_, err := c.
205+
Database("test").
206+
Collection("sessions").
207+
InsertOne(ctx, bson.M{
208+
"uid": "session-4",
209+
"authenticated": true,
210+
})
211+
if err != nil {
212+
return err
213+
}
214+
215+
_, err = c.
216+
Database("test").
217+
Collection("recorded_sessions").
218+
InsertOne(ctx, bson.M{
219+
"uid": "session-4",
220+
"message": "test output",
221+
"time": "2023-01-01T10:00:00Z",
222+
"width": 80,
223+
"height": 24,
224+
})
225+
if err != nil {
226+
return err
227+
}
228+
229+
migrates := migrate.NewMigrate(c.Database("test"), GenerateMigrations()[94])
230+
if err := migrates.Up(ctx, migrate.AllAvailable); err != nil {
231+
return err
232+
}
233+
234+
return nil
235+
},
236+
verify: func(tt *testing.T) {
237+
count, err := c.
238+
Database("test").
239+
Collection("recorded_sessions").
240+
CountDocuments(ctx, bson.M{"uid": "session-4"})
241+
require.NoError(tt, err)
242+
assert.Equal(tt, int64(1), count)
243+
},
244+
},
245+
}
246+
247+
for _, tc := range cases {
248+
t.Run(tc.description, func(tt *testing.T) {
249+
tt.Cleanup(func() {
250+
assert.NoError(tt, srv.Reset())
251+
})
252+
253+
require.NoError(tt, tc.setup())
254+
255+
migrates := migrate.NewMigrate(c.Database("test"), GenerateMigrations()[94])
256+
require.NoError(tt, migrates.Down(ctx, migrate.AllAvailable))
257+
258+
tc.verify(tt)
259+
})
260+
}
261+
}

0 commit comments

Comments
 (0)