expectation.go

package sqlmock

import (
	"database/sql/driver"
	"errors"
	"fmt"
)

// ExpectedQuery describes an expected query operation.
type ExpectedQuery struct {
	sql  string
	args []any
	rows *Rows
}

// ExpectQuery appends a query expectation to the mock queue.
func (m *Mock) ExpectQuery(sql string) *ExpectedQuery {
	e := &ExpectedQuery{sql: sql}
	m.mu.Lock()
	m.ops = append(m.ops, e)
	m.mu.Unlock()
	return e
}

// WithArgs sets expected arguments for the query expectation.
func (e *ExpectedQuery) WithArgs(args ...any) *ExpectedQuery {
	e.args = args
	return e
}

// WillReturnRows sets the rows returned when the query is matched.
func (e *ExpectedQuery) WillReturnRows(r *Rows) {
	e.rows = r
}

func (e *ExpectedQuery) match(op runtimeOp) error {
	if op.kind != "query" {
		return fmt.Errorf("expected QUERY, got %s", op.kind)
	}

	if e.sql != op.sql {
		return fmt.Errorf("query sql mismatch: %s", op.sql)
	}

	if !matchArgs(e.args, op.args) {
		return errors.New("query args mismatch")
	}

	return nil
}

// ExpectedExec describes an expected exec operation.
type ExpectedExec struct {
	sql    string
	args   []any
	result driver.Result
	err    error
}

// ExpectExec appends an exec expectation to the mock queue.
//
// By default, an exec expectation returns RowsAffected(1) unless overridden.
func (m *Mock) ExpectExec(sql string) *ExpectedExec {
	e := &ExpectedExec{
		sql:    sql,
		result: driver.RowsAffected(1),
	}

	m.mu.Lock()
	m.ops = append(m.ops, e)
	m.mu.Unlock()
	return e
}

// WithArgs sets expected arguments for the exec expectation.
func (e *ExpectedExec) WithArgs(args ...any) *ExpectedExec {
	e.args = args
	return e
}

// WillReturnResult sets the rows-affected result returned by the exec.
func (e *ExpectedExec) WillReturnResult(affected int64) {
	e.result = driver.RowsAffected(affected)
}

// WillReturnError configures the error returned by the exec expectation.
func (e *ExpectedExec) WillReturnError(err error) {
	e.err = err
}

func (e *ExpectedExec) match(op runtimeOp) error {
	if op.kind != "exec" {
		return fmt.Errorf("expected EXEC, got %s", op.kind)
	}

	if e.sql != op.sql {
		return fmt.Errorf("exec sql mismatch: %s", op.sql)
	}

	if !matchArgs(e.args, op.args) {
		return errors.New("exec args mismatch")
	}

	return e.err
}

// ExpectedBegin describes an expected begin-transaction operation.
type ExpectedBegin struct {
	err error
}

// ExpectBegin appends a begin-transaction expectation to the mock queue.
func (m *Mock) ExpectBegin() *ExpectedBegin {
	e := &ExpectedBegin{}
	m.mu.Lock()
	m.ops = append(m.ops, e)
	m.mu.Unlock()
	return e
}

// WillReturnError configures the error returned when begin is matched.
func (e *ExpectedBegin) WillReturnError(err error) {
	e.err = err
}

func (e *ExpectedBegin) match(op runtimeOp) error {
	if op.kind != "begin" {
		return fmt.Errorf("expected BEGIN, got %s", op.kind)
	}

	return e.err
}

// ExpectedCommit describes an expected commit operation.
type ExpectedCommit struct {
	err error
}

// ExpectCommit appends a commit expectation to the mock queue.
func (m *Mock) ExpectCommit() *ExpectedCommit {
	e := &ExpectedCommit{}
	m.mu.Lock()
	m.ops = append(m.ops, e)
	m.mu.Unlock()
	return e
}

// WillReturnError configures the error returned when commit is matched.
func (e *ExpectedCommit) WillReturnError(err error) {
	e.err = err
}

func (e *ExpectedCommit) match(op runtimeOp) error {
	if op.kind != "commit" {
		return fmt.Errorf("expected COMMIT, got %s", op.kind)
	}

	return e.err
}

// ExpectedRollback describes an expected rollback operation.
type ExpectedRollback struct {
	err error
}

// ExpectRollback appends a rollback expectation to the mock queue.
func (m *Mock) ExpectRollback() *ExpectedRollback {
	e := &ExpectedRollback{}
	m.mu.Lock()
	m.ops = append(m.ops, e)
	m.mu.Unlock()
	return e
}

// WillReturnError configures the error returned when rollback is matched.
func (e *ExpectedRollback) WillReturnError(err error) {
	e.err = err
}

func (e *ExpectedRollback) match(op runtimeOp) error {
	if op.kind != "rollback" {
		return fmt.Errorf("expected ROLLBACK, got %s", op.kind)
	}

	return e.err
}

// ExpectationsWereMet reports whether all queued expectations were consumed.
func (m *Mock) ExpectationsWereMet() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if len(m.ops) > 0 {
		return fmt.Errorf("not all expectations were met (%d remaining)", len(m.ops))
	}

	return nil
}