chore: making unit tests to try out various things

coder · Emyrk · Nov 1, 2024 · Oct 23, 2024 · Oct 23, 2024 · Oct 24, 2024
commit dce750db6cfc63592dd53c97b54315d9fc37687f
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)