Skip to content

Commit 11c20d6

Browse files
johnstcnjohnstcn
committed
use dbtestutil in httpmw test, refactor
1 parent c93782a commit 11c20d6

File tree

1 file changed

+69
-47
lines changed

1 file changed

+69
-47
lines changed

‎coderd/httpmw/workspaceagent_test.go

Lines changed: 69 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -10,35 +10,31 @@ import (
1010
"github.com/stretchr/testify/require"
1111

1212
"github.com/coder/coder/coderd/database"
13-
"github.com/coder/coder/coderd/database/dbfake"
1413
"github.com/coder/coder/coderd/database/dbgen"
14+
"github.com/coder/coder/coderd/database/dbtestutil"
1515
"github.com/coder/coder/coderd/httpmw"
1616
"github.com/coder/coder/codersdk"
1717
)
1818

1919
func TestWorkspaceAgent(t *testing.T) {
2020
t.Parallel()
2121

22-
setup := func(db database.Store, token uuid.UUID) *http.Request {
23-
r := httptest.NewRequest("GET", "/", nil)
24-
r.Header.Set(codersdk.SessionTokenHeader, token.String())
25-
return r
26-
}
27-
2822
t.Run("None", func(t *testing.T) {
2923
t.Parallel()
30-
db := dbfake.New()
31-
rtr := chi.NewRouter()
32-
rtr.Use(
33-
httpmw.ExtractWorkspaceAgent(httpmw.ExtractWorkspaceAgentConfig{
24+
db, pubsub := dbtestutil.NewDB(t)
25+
t.Cleanup(func() {
26+
_ = pubsub.Close()
27+
})
28+
29+
req, rtr := setup(t, db, uuid.New(), httpmw.ExtractWorkspaceAgent(
30+
httpmw.ExtractWorkspaceAgentConfig{
3431
DB: db,
3532
Optional: false,
36-
}),
37-
)
38-
rtr.Get("/", nil)
39-
r := setup(db, uuid.New())
33+
}))
34+
4035
rw := httptest.NewRecorder()
41-
rtr.ServeHTTP(rw, r)
36+
req.Header.Set(codersdk.SessionTokenHeader, uuid.New().String())
37+
rtr.ServeHTTP(rw, req)
4238

4339
res := rw.Result()
4440
defer res.Body.Close()
@@ -47,42 +43,68 @@ func TestWorkspaceAgent(t *testing.T) {
4743

4844
t.Run("Found", func(t *testing.T) {
4945
t.Parallel()
50-
db := dbfake.New()
51-
var (
52-
user = dbgen.User(t, db, database.User{})
53-
workspace = dbgen.Workspace(t, db, database.Workspace{
54-
OwnerID: user.ID,
55-
})
56-
job = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{})
57-
resource = dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
58-
JobID: job.ID,
59-
})
60-
_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
61-
WorkspaceID: workspace.ID,
62-
JobID: job.ID,
63-
})
64-
agent = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
65-
ResourceID: resource.ID,
66-
})
67-
)
68-
69-
rtr := chi.NewRouter()
70-
rtr.Use(
71-
httpmw.ExtractWorkspaceAgent(httpmw.ExtractWorkspaceAgentConfig{
46+
db, pubsub := dbtestutil.NewDB(t)
47+
t.Cleanup(func() {
48+
_ = pubsub.Close()
49+
})
50+
authToken := uuid.New()
51+
req, rtr := setup(t, db, authToken, httpmw.ExtractWorkspaceAgent(
52+
httpmw.ExtractWorkspaceAgentConfig{
7253
DB: db,
7354
Optional: false,
74-
}),
75-
)
76-
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
77-
_ = httpmw.WorkspaceAgent(r)
78-
rw.WriteHeader(http.StatusOK)
79-
})
80-
r := setup(db, agent.AuthToken)
55+
}))
56+
8157
rw := httptest.NewRecorder()
82-
rtr.ServeHTTP(rw, r)
58+
req.Header.Set(codersdk.SessionTokenHeader, authToken.String())
59+
rtr.ServeHTTP(rw, req)
8360

8461
res := rw.Result()
85-
defer res.Body.Close()
62+
t.Cleanup(func() { _ = res.Body.Close() })
8663
require.Equal(t, http.StatusOK, res.StatusCode)
8764
})
8865
}
66+
67+
func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Handler) http.Handler) (*http.Request, http.Handler) {
68+
t.Helper()
69+
org := dbgen.Organization(t, db, database.Organization{})
70+
user := dbgen.User(t, db, database.User{})
71+
templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
72+
OrganizationID: org.ID,
73+
CreatedBy: user.ID,
74+
})
75+
template := dbgen.Template(t, db, database.Template{
76+
OrganizationID: org.ID,
77+
ActiveVersionID: templateVersion.ID,
78+
CreatedBy: user.ID,
79+
})
80+
workspace := dbgen.Workspace(t, db, database.Workspace{
81+
OwnerID: user.ID,
82+
OrganizationID: org.ID,
83+
TemplateID: template.ID,
84+
})
85+
job := dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
86+
OrganizationID: org.ID,
87+
})
88+
resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
89+
JobID: job.ID,
90+
})
91+
_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
92+
WorkspaceID: workspace.ID,
93+
JobID: job.ID,
94+
TemplateVersionID: templateVersion.ID,
95+
})
96+
_ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
97+
ResourceID: resource.ID,
98+
AuthToken: authToken,
99+
})
100+
101+
req := httptest.NewRequest("GET", "/", nil)
102+
rtr := chi.NewRouter()
103+
rtr.Use(mw)
104+
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
105+
_ = httpmw.WorkspaceAgent(r)
106+
rw.WriteHeader(http.StatusOK)
107+
})
108+
109+
return req, rtr
110+
}

0 commit comments

Comments
 (0)