package dbtestutil

import (

"sync"

"testing"

"github.com/stretchr/testify/assert"

"golang.org/x/xerrors"

"github.com/coder/coder/v2/coderd/database"

)

type DBTx struct {

database.Store

mu sync.Mutex

done 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.

// Example:

//

// a := StartTx(t, db, opts)

// b := StartTx(t, db, opts)

//

// a.GetUsers(...)

// b.GetUsers(...)

//

// require.NoError(t, a.Done()

func StartTx(t *testing.T, db database.Store, opts *database.TxOptions) *DBTx {

done := 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 {

// If you recursively call InTx, then don't use this.

t.Logf("InTx called more than once: %d", count)

assert.NoError(t, xerrors.New("InTx called more than once, this is not allowed with the StartTx helper"))

}

<-done

// Just return nil. The caller should be checking their own errors.

return nil

}, opts)

finalErr <- err

}()

txStore := <-txC

close(txC)

return &DBTx{Store: txStore, done: done, finalErr: finalErr}

}

// 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()

close(tx.done)

return <-tx.finalErr

}