chore: making unit tests to try out various things

Emyrk · Emyrk · commit dce750db6cfc · 2024-10-28T17:16:05.000-04:00
diff --git a/coderd/database/dbtestutil/tx.go b/coderd/database/dbtestutil/tx.go
@@ -0,0 +1,65 @@
+package dbtestutil
+
+import (
+	"database/sql"
+	"sync"
+	"testing"
+
+	"github.com/coder/coder/v2/coderd/database"
+)
+
+type DBTx struct {
+	database.Store
+	mu       sync.Mutex
+	err      error
+	errC     chan error
+	finalErr chan error
+}
+
+// StartTx starts a transaction and returns a DBTx object. This allows running
+// 2 transactions concurrently in a test more easily.
+func StartTx(t *testing.T, db database.Store, opts *sql.TxOptions) *DBTx {
+	errC := make(chan error)
+	finalErr := make(chan error)
+	txC := make(chan database.Store)
+
+	go func() {
+		t.Helper()
+		once := sync.Once{}
+		count := 0
+
+		err := db.InTx(func(store database.Store) error {
+			// InTx can be retried
+			once.Do(func() {
+				txC <- store
+			})
+			count++
+			if count > 1 {
+				t.Logf("InTx called more than once: %d", count)
+			}
+			return <-errC
+		}, opts)
+		finalErr <- err
+	}()
+
+	txStore := <-txC
+	close(txC)
+
+	return &DBTx{Store: txStore, errC: errC, finalErr: finalErr}
+}
+
+func (tx *DBTx) SetError(err error) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.err = err
+}
+
+// Done can only be called once. If you call it twice, it will panic.
+func (tx *DBTx) Done() error {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+
+	tx.errC <- tx.err
+	close(tx.errC)
+	return <-tx.finalErr
+}
diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -701,7 +701,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 
 		if initial, changed, enabled := featureChanged(codersdk.FeatureTemplateRBAC); shouldUpdate(initial, changed, enabled) {
 			if enabled {
-				committer := committer{
+				committer := Committer{
 					Log:      api.Logger.Named("quota_committer"),
 					Database: api.Database,
 				}
diff --git a/enterprise/coderd/workspacequota.go b/enterprise/coderd/workspacequota.go
@@ -18,12 +18,12 @@ import (
 	"github.com/coder/coder/v2/provisionerd/proto"
 )
 
-type committer struct {
+type Committer struct {
 	Log      slog.Logger
 	Database database.Store
 }
 
-func (c *committer) CommitQuota(
+func (c *Committer) CommitQuota(
 	ctx context.Context, request *proto.CommitQuotaRequest,
 ) (*proto.CommitQuotaResponse, error) {
 	jobID, err := uuid.Parse(request.JobId)
diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go

Original file line number	Diff line number	Diff line change
`@@ -701,7 +701,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {`
`701`	`701`
`702`	`702`	`if initial, changed, enabled := featureChanged(codersdk.FeatureTemplateRBAC); shouldUpdate(initial, changed, enabled) {`
`703`	`703`	`if enabled {`
`704`		`- committer := committer{`
	`704`	`+ committer := Committer{`
`705`	`705`	`Log: api.Logger.Named("quota_committer"),`
`706`	`706`	`Database: api.Database,`
`707`	`707`	`}`
Original file line number	Diff line number	Diff line change
`@@ -18,12 +18,12 @@ import (`
`18`	`18`	`"github.com/coder/coder/v2/provisionerd/proto"`
`19`	`19`	`)`
`20`	`20`
`21`		`-type committer struct {`
	`21`	`+type Committer struct {`
`22`	`22`	`Log slog.Logger`
`23`	`23`	`Database database.Store`
`24`	`24`	`}`
`25`	`25`
`26`		`-func (c *committer) CommitQuota(`
	`26`	`+func (c *Committer) CommitQuota(`
`27`	`27`	`ctx context.Context, request *proto.CommitQuotaRequest,`
`28`	`28`	`) (*proto.CommitQuotaResponse, error) {`
`29`	`29`	`jobID, err := uuid.Parse(request.JobId)`