Skip to content

Commit 1610f34

Browse files
easyCZeasyCZ
committed
fix up
1 parent efc9c79 commit 1610f34

File tree

5 files changed

+103
-105
lines changed

5 files changed

+103
-105
lines changed

‎components/usage/pkg/controller/reconciler_test.go

Lines changed: 45 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -15,93 +15,58 @@ import (
1515
)
1616

1717
func TestUsageReconciler_ReconcileTimeRange(t *testing.T) {
18-
type Scenario struct {
19-
NowFunc func() time.Time
20-
21-
Instances []db.WorkspaceInstance
22-
23-
Expected *UsageReconcileStatus
24-
ExpectedReportSize int
25-
ExpectedAttributedRuntime map[string]int64
26-
}
27-
2818
startOfMay := time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC)
2919
startOfJune := time.Date(2022, 06, 1, 0, 00, 00, 00, time.UTC)
3020

31-
scenarios := []struct {
32-
Name string
33-
Setup func(t *testing.T) Scenario
34-
}{
35-
{
36-
Name: "3 workspace instances, one invalid",
37-
Setup: func(t *testing.T) Scenario {
38-
teamID := uuid.New()
39-
scenarioRunTime := time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC)
40-
41-
instances := []db.WorkspaceInstance{
42-
// Ran throughout the reconcile period
43-
dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{
44-
ID: uuid.New(),
45-
UsageAttributionID: db.NewTeamAttributionID(teamID.String()),
46-
CreationTime: db.NewVarcharTime(time.Date(2022, 05, 1, 00, 00, 00, 00, time.UTC)),
47-
StartedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 00, 00, 00, 00, time.UTC)),
48-
StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)),
49-
}),
50-
// Still running
51-
dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{
52-
ID: uuid.New(),
53-
UsageAttributionID: db.NewTeamAttributionID(teamID.String()),
54-
CreationTime: db.NewVarcharTime(time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC)),
55-
StartedTime: db.NewVarcharTime(time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC)),
56-
}),
57-
// No creation time, invalid record
58-
dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{
59-
ID: uuid.New(),
60-
UsageAttributionID: db.NewTeamAttributionID(teamID.String()),
61-
StartedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)),
62-
StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)),
63-
}),
64-
}
21+
teamID := uuid.New()
22+
scenarioRunTime := time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC)
6523

66-
expectedRuntime := instances[0].WorkspaceRuntimeSeconds(scenarioRunTime) + instances[1].WorkspaceRuntimeSeconds(scenarioRunTime)
67-
68-
return Scenario{
69-
Instances: instances,
70-
NowFunc: func() time.Time { return scenarioRunTime },
71-
ExpectedReportSize: 1,
72-
ExpectedAttributedRuntime: map[string]int64{
73-
teamID.String(): int64(expectedRuntime),
74-
},
75-
Expected: &UsageReconcileStatus{
76-
StartTime: startOfMay,
77-
EndTime: startOfJune,
78-
WorkspaceInstances: 2,
79-
InvalidWorkspaceInstances: 1,
80-
},
81-
}
82-
},
83-
},
24+
instances := []db.WorkspaceInstance{
25+
// Ran throughout the reconcile period
26+
dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{
27+
ID: uuid.New(),
28+
UsageAttributionID: db.NewTeamAttributionID(teamID.String()),
29+
CreationTime: db.NewVarcharTime(time.Date(2022, 05, 1, 00, 00, 00, 00, time.UTC)),
30+
StartedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 00, 00, 00, 00, time.UTC)),
31+
StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)),
32+
}),
33+
// Still running
34+
dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{
35+
ID: uuid.New(),
36+
UsageAttributionID: db.NewTeamAttributionID(teamID.String()),
37+
CreationTime: db.NewVarcharTime(time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC)),
38+
StartedTime: db.NewVarcharTime(time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC)),
39+
}),
40+
// No creation time, invalid record
41+
dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{
42+
ID: uuid.New(),
43+
UsageAttributionID: db.NewTeamAttributionID(teamID.String()),
44+
StartedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)),
45+
StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)),
46+
}),
8447
}
8548

86-
for _, s := range scenarios {
87-
t.Run(s.Name, func(t *testing.T) {
88-
scenario := s.Setup(t)
89-
90-
conn := db.ConnectForTests(t)
91-
require.NoError(t, conn.Create(scenario.Instances).Error)
49+
expectedRuntime := instances[0].WorkspaceRuntimeSeconds(scenarioRunTime) + instances[1].WorkspaceRuntimeSeconds(scenarioRunTime)
9250

93-
reconciler := &UsageReconciler{
94-
billingController: &NoOpBillingController{},
95-
nowFunc: scenario.NowFunc,
96-
conn: conn,
97-
}
98-
status, report, err := reconciler.ReconcileTimeRange(context.Background(), startOfMay, startOfJune)
99-
require.NoError(t, err)
51+
conn := db.ConnectForTests(t)
52+
dbtest.CreateWorkspaceInstances(t, conn, instances...)
10053

101-
require.Len(t, report, scenario.ExpectedReportSize)
102-
require.Equal(t, scenario.Expected, status)
103-
require.Equal(t, scenario.ExpectedAttributedRuntime, report.RuntimeSummaryForTeams(scenario.NowFunc()))
104-
105-
})
54+
reconciler := &UsageReconciler{
55+
billingController: &NoOpBillingController{},
56+
nowFunc: func() time.Time { return scenarioRunTime },
57+
conn: conn,
10658
}
59+
status, report, err := reconciler.ReconcileTimeRange(context.Background(), startOfMay, startOfJune)
60+
require.NoError(t, err)
61+
62+
require.Len(t, report, 1)
63+
require.Equal(t, &UsageReconcileStatus{
64+
StartTime: startOfMay,
65+
EndTime: startOfJune,
66+
WorkspaceInstances: 2,
67+
InvalidWorkspaceInstances: 1,
68+
}, status)
69+
require.Equal(t, map[string]int64{
70+
teamID.String(): int64(expectedRuntime),
71+
}, report.RuntimeSummaryForTeams(scenarioRunTime))
10772
}

‎components/usage/pkg/db/conn.go

Lines changed: 19 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -72,25 +72,25 @@ func ConnectForTests(t *testing.T) *gorm.DB {
7272
})
7373
require.NoError(t, err, "Failed to establish connection to DB. In a workspace, run `leeway build components/usage:init-testdb` once to bootstrap the DB.")
7474

75-
t.Cleanup(func() {
76-
// Delete records for known models from the DB
77-
log.Info("Cleaning up DB between connections.")
78-
for _, model := range []interface{}{
79-
&WorkspaceInstance{},
80-
&Workspace{},
81-
&Project{},
82-
&Team{},
83-
&TeamMembership{},
84-
} {
85-
// See https://gorm.io/docs/delete.html#Block-Global-Delete
86-
tx := conn.Where("1 = 1").Delete(model)
87-
require.NoError(t, tx.Error)
88-
}
89-
90-
rawConn, err := conn.DB()
91-
require.NoError(t, err)
92-
require.NoError(t, rawConn.Close(), "must close database connection")
93-
})
75+
//t.Cleanup(func() {
76+
// // Delete records for known models from the DB
77+
// log.Debug("Cleaning up DB between connections.")
78+
// for _, model := range []interface{}{
79+
// &WorkspaceInstance{},
80+
// &Workspace{},
81+
// &Project{},
82+
// &Team{},
83+
// &TeamMembership{},
84+
// } {
85+
// // See https://gorm.io/docs/delete.html#Block-Global-Delete
86+
// tx := conn.Where("1 = 1").Delete(model)
87+
// require.NoError(t, tx.Error)
88+
// }
89+
//
90+
// rawConn, err := conn.DB()
91+
// require.NoError(t, err)
92+
// require.NoError(t, rawConn.Close(), "must close database connection")
93+
//})
9494

9595
return conn
9696
}

‎components/usage/pkg/db/dbtest/workspace.go

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import (
88
"fmt"
99
"github.com/gitpod-io/gitpod/usage/pkg/db"
1010
"github.com/google/uuid"
11+
"github.com/stretchr/testify/require"
12+
"gorm.io/gorm"
1113
"math/rand"
1214
"testing"
1315
)
@@ -76,3 +78,16 @@ func randSeq(n int) string {
7678
}
7779
return string(b)
7880
}
81+
82+
func CreateWorkspace(t *testing.T, conn *gorm.DB, workspace db.Workspace) db.Workspace {
83+
t.Helper()
84+
85+
ws := &workspace
86+
require.NoError(t, conn.Create(&ws).Error)
87+
88+
t.Cleanup(func() {
89+
require.NoError(t, conn.Where("id = ?", workspace.ID).Delete(ws).Error)
90+
})
91+
92+
return *ws
93+
}

‎components/usage/pkg/db/dbtest/workspace_instance.go

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import (
88
"database/sql"
99
"github.com/gitpod-io/gitpod/usage/pkg/db"
1010
"github.com/google/uuid"
11+
"github.com/stretchr/testify/require"
12+
"gorm.io/gorm"
1113
"testing"
1214
"time"
1315
)
@@ -87,3 +89,23 @@ func NewWorkspaceInstance(t *testing.T, instance db.WorkspaceInstance) db.Worksp
8789
PhasePersisted: "",
8890
}
8991
}
92+
93+
func CreateWorkspaceInstances(t *testing.T, conn *gorm.DB, instances ...db.WorkspaceInstance) []db.WorkspaceInstance {
94+
t.Helper()
95+
96+
var records []db.WorkspaceInstance
97+
var ids []string
98+
for _, instance := range instances {
99+
record := NewWorkspaceInstance(t, instance)
100+
records = append(records, record)
101+
ids = append(ids, record.ID.String())
102+
}
103+
104+
require.NoError(t, conn.CreateInBatches(&records, 1000).Error)
105+
106+
t.Cleanup(func() {
107+
require.NoError(t, conn.Where(ids).Delete(&db.WorkspaceInstance{}).Error)
108+
})
109+
110+
return records
111+
}

‎components/usage/pkg/db/workspace_instance_test.go

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -179,10 +179,7 @@ func TestListWorkspaceInstancesInRange(t *testing.T) {
179179
all = append(all, valid...)
180180
all = append(all, invalid...)
181181

182-
for _, instance := range all {
183-
tx := conn.Create(&instance)
184-
require.NoError(t, tx.Error)
185-
}
182+
dbtest.CreateWorkspaceInstances(t, conn, all...)
186183

187184
startOfMay := time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC)
188185
startOfJune := time.Date(2022, 06, 1, 0, 00, 00, 00, time.UTC)
@@ -208,8 +205,7 @@ func TestListWorkspaceInstancesInRange_InBatches(t *testing.T) {
208205

209206
}
210207

211-
tx := conn.CreateInBatches(&instances, 1000)
212-
require.NoError(t, tx.Error)
208+
dbtest.CreateWorkspaceInstances(t, conn, instances...)
213209

214210
startOfMay := time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC)
215211
startOfJune := time.Date(2022, 06, 1, 0, 00, 00, 00, time.UTC)

0 commit comments

Comments
 (0)