Skip to content

Commit 50ea251

Browse files
EmyrkEmyrk
committed
feat: Add database data generator to make fakedbs easier to populate
One example shown in groupparm_test.go. More examples coming
1 parent f4d6afb commit 50ea251

File tree

2 files changed

+214
-17
lines changed

2 files changed

+214
-17
lines changed

‎coderd/database/databasefake/generator.go

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,212 @@
1+
package databasefake
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/coder/coder/coderd/database"
9+
"github.com/google/uuid"
10+
"github.com/moby/moby/pkg/namesgenerator"
11+
"github.com/stretchr/testify/require"
12+
)
13+
14+
const primaryOrgName = "primary-org"
15+
16+
type Generator struct {
17+
// names is a map of names to uuids.
18+
names map[string]uuid.UUID
19+
primaryOrg *database.Organization
20+
testT *testing.T
21+
22+
db database.Store
23+
}
24+
25+
func NewGenerator(t *testing.T, db database.Store) *Generator {
26+
if _, ok := db.(FakeDatabase); !ok {
27+
panic("Generator db must be a FakeDatabase")
28+
}
29+
return &Generator{
30+
names: make(map[string]uuid.UUID),
31+
testT: t,
32+
db: db,
33+
}
34+
}
35+
36+
// PrimaryOrg is to keep all resources in the same default org if not
37+
// specified.
38+
func (g *Generator) PrimaryOrg(ctx context.Context) database.Organization {
39+
if g.primaryOrg == nil {
40+
org := g.Organization(ctx, "primary-org", database.Organization{
41+
ID: g.Lookup(primaryOrgName),
42+
Name: primaryOrgName,
43+
Description: "This is the default primary organization for all tests",
44+
CreatedAt: time.Now(),
45+
UpdatedAt: time.Now(),
46+
})
47+
g.primaryOrg = &org
48+
}
49+
50+
return *g.primaryOrg
51+
}
52+
53+
func populate[DBType any](ctx context.Context, g *Generator, name string, seed DBType) DBType {
54+
out := g.Populate(ctx, map[string]interface{}{
55+
name: seed,
56+
})
57+
return out[name].(DBType)
58+
}
59+
60+
func (g *Generator) Group(ctx context.Context, name string, seed database.Group) database.Group {
61+
return populate(ctx, g, name, seed)
62+
}
63+
64+
func (g *Generator) Organization(ctx context.Context, name string, seed database.Organization) database.Organization {
65+
return populate(ctx, g, name, seed)
66+
}
67+
68+
func (g *Generator) Workspace(ctx context.Context, name string, seed database.Workspace) database.Workspace {
69+
return populate(ctx, g, name, seed)
70+
}
71+
72+
func (g *Generator) Template(ctx context.Context, name string, seed database.Template) database.Template {
73+
return populate(ctx, g, name, seed)
74+
}
75+
76+
func (g *Generator) TemplateVersion(ctx context.Context, name string, seed database.TemplateVersion) database.TemplateVersion {
77+
return populate(ctx, g, name, seed)
78+
}
79+
80+
func (g *Generator) WorkspaceBuild(ctx context.Context, name string, seed database.WorkspaceBuild) database.WorkspaceBuild {
81+
return populate(ctx, g, name, seed)
82+
}
83+
84+
func (g *Generator) User(ctx context.Context, name string, seed database.User) database.User {
85+
return populate(ctx, g, name, seed)
86+
}
87+
88+
func (g *Generator) Populate(ctx context.Context, seed map[string]interface{}) map[string]interface{} {
89+
db := g.db
90+
t := g.testT
91+
92+
for name, v := range seed {
93+
switch orig := v.(type) {
94+
case database.Template:
95+
template, err := db.InsertTemplate(ctx, database.InsertTemplateParams{
96+
ID: g.Lookup(name),
97+
CreatedAt: time.Now(),
98+
UpdatedAt: time.Now(),
99+
OrganizationID: takeFirst(orig.OrganizationID, g.PrimaryOrg(ctx).ID),
100+
Name: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
101+
Provisioner: takeFirst(orig.Provisioner, database.ProvisionerTypeEcho),
102+
ActiveVersionID: takeFirst(orig.ActiveVersionID, uuid.New()),
103+
Description: takeFirst(orig.Description, namesgenerator.GetRandomName(1)),
104+
DefaultTTL: takeFirst(orig.DefaultTTL, 3600),
105+
CreatedBy: takeFirst(orig.CreatedBy, uuid.New()),
106+
Icon: takeFirst(orig.Icon, namesgenerator.GetRandomName(1)),
107+
UserACL: orig.UserACL,
108+
GroupACL: orig.GroupACL,
109+
DisplayName: takeFirst(orig.DisplayName, namesgenerator.GetRandomName(1)),
110+
AllowUserCancelWorkspaceJobs: takeFirst(orig.AllowUserCancelWorkspaceJobs, true),
111+
})
112+
require.NoError(t, err, "insert template")
113+
114+
seed[name] = template
115+
case database.Workspace:
116+
workspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
117+
ID: g.Lookup(name),
118+
CreatedAt: time.Now(),
119+
UpdatedAt: time.Now(),
120+
OrganizationID: takeFirst(orig.OrganizationID, g.PrimaryOrg(ctx).ID),
121+
TemplateID: takeFirst(orig.TemplateID, uuid.New()),
122+
Name: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
123+
AutostartSchedule: orig.AutostartSchedule,
124+
Ttl: orig.Ttl,
125+
})
126+
require.NoError(t, err, "insert workspace")
127+
128+
seed[name] = workspace
129+
case database.WorkspaceBuild:
130+
build, err := db.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{
131+
ID: g.Lookup(name),
132+
CreatedAt: time.Now(),
133+
UpdatedAt: time.Now(),
134+
WorkspaceID: takeFirst(orig.WorkspaceID, uuid.New()),
135+
TemplateVersionID: takeFirst(orig.TemplateVersionID, uuid.New()),
136+
BuildNumber: takeFirst(orig.BuildNumber, 0),
137+
Transition: takeFirst(orig.Transition, database.WorkspaceTransitionStart),
138+
InitiatorID: takeFirst(orig.InitiatorID, uuid.New()),
139+
JobID: takeFirst(orig.InitiatorID, uuid.New()),
140+
ProvisionerState: []byte{},
141+
Deadline: time.Now(),
142+
Reason: takeFirst(orig.Reason, database.BuildReasonInitiator),
143+
})
144+
require.NoError(t, err, "insert workspace build")
145+
146+
seed[name] = build
147+
case database.User:
148+
user, err := db.InsertUser(ctx, database.InsertUserParams{
149+
ID: g.Lookup(name),
150+
Email: takeFirst(orig.Email, namesgenerator.GetRandomName(1)),
151+
Username: takeFirst(orig.Username, namesgenerator.GetRandomName(1)),
152+
HashedPassword: []byte{},
153+
CreatedAt: time.Now(),
154+
UpdatedAt: time.Now(),
155+
RBACRoles: []string{},
156+
LoginType: takeFirst(orig.LoginType, database.LoginTypePassword),
157+
})
158+
require.NoError(t, err, "insert user")
159+
160+
seed[name] = user
161+
162+
case database.Organization:
163+
org, err := db.InsertOrganization(ctx, database.InsertOrganizationParams{
164+
ID: g.Lookup(name),
165+
Name: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
166+
Description: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
167+
CreatedAt: time.Now(),
168+
UpdatedAt: time.Now(),
169+
})
170+
require.NoError(t, err, "insert organization")
171+
172+
seed[name] = org
173+
174+
case database.Group:
175+
org, err := db.InsertGroup(ctx, database.InsertGroupParams{
176+
ID: g.Lookup(name),
177+
Name: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
178+
OrganizationID: takeFirst(orig.OrganizationID, g.PrimaryOrg(ctx).ID),
179+
AvatarURL: takeFirst(orig.Name, "https://logo.example.com"),
180+
QuotaAllowance: takeFirst(orig.QuotaAllowance, 0),
181+
})
182+
require.NoError(t, err, "insert organization")
183+
184+
seed[name] = org
185+
}
186+
}
187+
return seed
188+
}
189+
190+
func (tc *Generator) Lookup(name string) uuid.UUID {
191+
if tc.names == nil {
192+
tc.names = make(map[string]uuid.UUID)
193+
}
194+
if id, ok := tc.names[name]; ok {
195+
return id
196+
}
197+
id := uuid.New()
198+
tc.names[name] = id
199+
return id
200+
}
201+
202+
// takeFirst will take the first non-empty value.
203+
func takeFirst[Value comparable](values ...Value) Value {
204+
var empty Value
205+
for _, v := range values {
206+
if v != empty {
207+
return v
208+
}
209+
}
210+
// If all empty, return empty
211+
return empty
212+
}

‎coderd/httpmw/groupparam_test.go

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -24,23 +24,8 @@ func TestGroupParam(t *testing.T) {
2424

2525
ctx, _ := testutil.Context(t)
2626
db := databasefake.New()
27-
28-
orgID := uuid.New()
29-
organization, err := db.InsertOrganization(ctx, database.InsertOrganizationParams{
30-
ID: orgID,
31-
Name: "banana",
32-
Description: "wowie",
33-
CreatedAt: database.Now(),
34-
UpdatedAt: database.Now(),
35-
})
36-
require.NoError(t, err)
37-
38-
group, err := db.InsertGroup(ctx, database.InsertGroupParams{
39-
ID: uuid.New(),
40-
Name: "yeww",
41-
OrganizationID: organization.ID,
42-
})
43-
require.NoError(t, err)
27+
gen := databasefake.NewGenerator(t, db)
28+
group := gen.Group(ctx, "group", database.Group{})
4429

4530
return db, group
4631
}

0 commit comments

Comments
 (0)