pluginsdk/plugin/wasmguest/bnwasm/driver_test.go

package bnwasm

import (
	"context"
	"database/sql"
	"errors"
	"reflect"
	"slices"
	"testing"
	"time"

	abiv1 "git.dev.alexdunmow.com/block/pluginsdk/abi/v1"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"google.golang.org/protobuf/proto"
)

// recordedCall captures one db.* host call for ordered-sequence assertions.
type recordedCall struct {
	method   string
	sql      string
	args     []*abiv1.DbValue
	txHandle uint64
}

// fakeHost is a scripted db.* transport: it records every call and answers from
// programmable fields, so the guest driver is exercised end-to-end with no real
// Postgres and the exact host-call sequence can be asserted.
type fakeHost struct {
	calls []recordedCall

	queryResp *abiv1.DbRowsResponse // returned for db.query
	execRows  int64                 // rows_affected for db.exec
	txHandle  uint64                // handle handed out on db.tx_begin
	dbError   *abiv1.DbError        // if set, attached to the next query/exec response
	failWith  error                 // if set, the transport returns this transport-level error
}

func (f *fakeHost) call(method string, req, resp proto.Message) error {
	rc := recordedCall{method: method}
	switch r := req.(type) {
	case *abiv1.DbQueryRequest:
		rc.sql, rc.args, rc.txHandle = r.GetSql(), r.GetArgs(), r.GetTxHandle()
	case *abiv1.DbExecRequest:
		rc.sql, rc.args, rc.txHandle = r.GetSql(), r.GetArgs(), r.GetTxHandle()
	case *abiv1.DbTxCommitRequest:
		rc.txHandle = r.GetTxHandle()
	case *abiv1.DbTxRollbackRequest:
		rc.txHandle = r.GetTxHandle()
	}
	f.calls = append(f.calls, rc)

	if f.failWith != nil {
		return f.failWith
	}
	switch out := resp.(type) {
	case *abiv1.DbRowsResponse:
		if f.queryResp != nil {
			proto.Merge(out, f.queryResp)
		}
		out.Error = f.dbError
	case *abiv1.DbExecResponse:
		out.RowsAffected = f.execRows
		out.Error = f.dbError
	case *abiv1.DbTxBeginResponse:
		out.TxHandle = f.txHandle
		out.Error = f.dbError
	case *abiv1.DbTxCommitResponse:
		out.Error = f.dbError
	case *abiv1.DbTxRollbackResponse:
		out.Error = f.dbError
	}
	return nil
}

func (f *fakeHost) methods() []string {
	ms := make([]string, len(f.calls))
	for i, c := range f.calls {
		ms[i] = c.method
	}
	return ms
}

func rowsWith(columns []string, cells ...*abiv1.DbValue) *abiv1.DbRowsResponse {
	return &abiv1.DbRowsResponse{Columns: columns, Rows: []*abiv1.DbRow{{Values: cells}}}
}

// deref returns the value a NewDest()-style pointer points at.
func deref(ptr any) any { return reflect.ValueOf(ptr).Elem().Interface() }

func equalScan(want, got any) bool {
	if w, ok := want.(time.Time); ok {
		g, ok := got.(time.Time)
		return ok && w.Equal(g)
	}
	return reflect.DeepEqual(want, got)
}

// TestDbValueScanRoundTrip proves every DbValue variant round-trips through a
// SELECT and scans into its canonical Go type — the core contract WO-WZ-007
// mirrors. Driven by the shared DbValueFixtures table.
func TestDbValueScanRoundTrip(t *testing.T) {
	for _, fx := range DbValueFixtures {
		t.Run(fx.Name, func(t *testing.T) {
			host := &fakeHost{queryResp: rowsWith([]string{fx.Name}, fx.Value)}
			pool := NewPool(host.call)

			dest := fx.NewDest()
			if err := pool.QueryRow(context.Background(), "SELECT x").Scan(dest); err != nil {
				t.Fatalf("scan %s: %v", fx.Name, err)
			}
			if got := deref(dest); !equalScan(fx.Want, got) {
				t.Fatalf("scan %s: want %#v, got %#v", fx.Name, fx.Want, got)
			}
		})
	}
}

// TestDbValueDriverValue proves the database/sql "bnwasm" driver maps every
// DbValue variant to the documented driver.Value.
func TestDbValueDriverValue(t *testing.T) {
	for _, fx := range DbValueFixtures {
		t.Run(fx.Name, func(t *testing.T) {
			if got := driverValue(fx.Value); !equalScan(fx.DriverValue, got) {
				t.Fatalf("driverValue %s: want %#v, got %#v", fx.Name, fx.DriverValue, got)
			}
		})
	}
}

func TestExecRowsAffected(t *testing.T) {
	host := &fakeHost{execRows: 5}
	pool := NewPool(host.call)

	tag, err := pool.Exec(context.Background(), "UPDATE t SET x = 1")
	if err != nil {
		t.Fatal(err)
	}
	if tag.RowsAffected() != 5 {
		t.Fatalf("rows affected: want 5, got %d", tag.RowsAffected())
	}
	if host.calls[0].txHandle != 0 {
		t.Fatalf("autocommit exec should carry tx_handle 0, got %d", host.calls[0].txHandle)
	}
}

// TestTxCommitSequence asserts the ordered host calls for a begin/exec/commit
// transaction, and that every statement inside carries the tx handle.
func TestTxCommitSequence(t *testing.T) {
	host := &fakeHost{txHandle: 77, execRows: 1}
	pool := NewPool(host.call)
	ctx := context.Background()

	tx, err := pool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := tx.Exec(ctx, "INSERT INTO t VALUES ($1)", 1); err != nil {
		t.Fatal(err)
	}
	if err := tx.Commit(ctx); err != nil {
		t.Fatal(err)
	}

	wantMethods := []string{methodTxBegin, methodExec, methodTxCommit}
	if got := host.methods(); !reflect.DeepEqual(got, wantMethods) {
		t.Fatalf("method sequence: want %v, got %v", wantMethods, got)
	}
	if host.calls[1].txHandle != 77 {
		t.Fatalf("in-tx exec should carry handle 77, got %d", host.calls[1].txHandle)
	}
	if host.calls[2].txHandle != 77 {
		t.Fatalf("commit should carry handle 77, got %d", host.calls[2].txHandle)
	}
}

// TestTxRollbackThenAutocommit asserts a rollback sequence, and that a query
// issued afterward (on the pool) carries NO tx handle.
func TestTxRollbackThenAutocommit(t *testing.T) {
	host := &fakeHost{txHandle: 9, execRows: 1}
	pool := NewPool(host.call)
	ctx := context.Background()

	tx, err := pool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if err := tx.Rollback(ctx); err != nil {
		t.Fatal(err)
	}
	if _, err := pool.Exec(ctx, "SELECT 1"); err != nil {
		t.Fatal(err)
	}

	wantMethods := []string{methodTxBegin, methodTxRollback, methodExec}
	if got := host.methods(); !reflect.DeepEqual(got, wantMethods) {
		t.Fatalf("method sequence: want %v, got %v", wantMethods, got)
	}
	if h := host.calls[1].txHandle; h != 9 {
		t.Fatalf("rollback should carry handle 9, got %d", h)
	}
	if h := host.calls[2].txHandle; h != 0 {
		t.Fatalf("post-rollback exec must carry tx_handle 0, got %d", h)
	}
}

func TestTxUseAfterCloseRejected(t *testing.T) {
	host := &fakeHost{txHandle: 3}
	tx, err := NewPool(host.call).Begin(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if err := tx.Commit(context.Background()); err != nil {
		t.Fatal(err)
	}
	if _, err := tx.Exec(context.Background(), "SELECT 1"); !errors.Is(err, pgx.ErrTxClosed) {
		t.Fatalf("exec after commit: want ErrTxClosed, got %v", err)
	}
	if err := tx.Rollback(context.Background()); !errors.Is(err, pgx.ErrTxClosed) {
		t.Fatalf("rollback after commit: want ErrTxClosed, got %v", err)
	}
}

func TestNestedTxRejected(t *testing.T) {
	host := &fakeHost{txHandle: 1}
	tx, err := NewPool(host.call).Begin(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if _, err := tx.Begin(context.Background()); !errors.Is(err, errNestedTx) {
		t.Fatalf("nested Begin: want errNestedTx, got %v", err)
	}
}

func TestNamedArgsRejected(t *testing.T) {
	host := &fakeHost{}
	pool := NewPool(host.call)
	_, err := pool.Query(context.Background(), "SELECT $1", pgx.NamedArgs{"a": 1})
	if err == nil {
		t.Fatal("expected named-args rejection, got nil")
	}
}

// TestDbErrorSurfacesAsPgError proves a DbError carrying a SQLSTATE reaches the
// plugin as a *pgconn.PgError, so errors.As-based constraint checks keep working.
func TestDbErrorSurfacesAsPgError(t *testing.T) {
	host := &fakeHost{dbError: &abiv1.DbError{Code: "23505", Message: "duplicate key"}}
	_, err := NewPool(host.call).Exec(context.Background(), "INSERT INTO t VALUES (1)")
	var pgErr *pgconn.PgError
	if !errors.As(err, &pgErr) {
		t.Fatalf("want *pgconn.PgError, got %T (%v)", err, err)
	}
	if pgErr.Code != "23505" {
		t.Fatalf("want SQLSTATE 23505, got %q", pgErr.Code)
	}
}

func TestNoHostFailsCleanly(t *testing.T) {
	pool := NewPool(nil)
	if _, err := pool.Exec(context.Background(), "SELECT 1"); !errors.Is(err, errNoHost) {
		t.Fatalf("want errNoHost, got %v", err)
	}
}

func TestQueryRowNoRows(t *testing.T) {
	host := &fakeHost{queryResp: &abiv1.DbRowsResponse{Columns: []string{"x"}}}
	var x int
	err := NewPool(host.call).QueryRow(context.Background(), "SELECT x").Scan(&x)
	if !errors.Is(err, pgx.ErrNoRows) {
		t.Fatalf("want ErrNoRows, got %v", err)
	}
}

// TestDatabaseSQLDriverRegistered proves the "bnwasm" driver is registered and
// drives a SELECT end-to-end through database/sql against the fake host.
func TestDatabaseSQLDriverRegistered(t *testing.T) {
	host := &fakeHost{queryResp: rowsWith(
		[]string{"id", "name"},
		&abiv1.DbValue{Kind: &abiv1.DbValue_Int64Value{Int64Value: 1}},
		&abiv1.DbValue{Kind: &abiv1.DbValue_StringValue{StringValue: "alice"}},
	)}
	db := sql.OpenDB(NewConnector(host.call))
	defer func() { _ = db.Close() }()

	var id int
	var name string
	if err := db.QueryRowContext(context.Background(), "SELECT id, name FROM t").Scan(&id, &name); err != nil {
		t.Fatal(err)
	}
	if id != 1 || name != "alice" {
		t.Fatalf("got id=%d name=%q", id, name)
	}
	if !slicesContains(sql.Drivers(), DriverName) {
		t.Fatalf("driver %q not registered; have %v", DriverName, sql.Drivers())
	}
}

func slicesContains(s []string, v string) bool {
	return slices.Contains(s, v)
}