test: fix api unit tests

Signed-off-by: Elias Van Ootegem <elias@vega.xyz>

datanode/api/trading_test.go

@@ -18,7 +18,6 @@ package api_test
 import (
 	"context"
 	"fmt"
-	"io"
 	"net"
 	"testing"
 	"time"
@@ -40,8 +39,6 @@ import (
 
 	"github.com/golang/mock/gomock"
 	"github.com/golang/protobuf/proto"
-	"github.com/jackc/pgconn"
-	"github.com/jackc/pgx/v4"
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc"
@@ -98,9 +95,8 @@ func getTestGRPCServer(t *testing.T, ctx context.Context) (tidy func(), conn *gr
 
 	conf.CandlesV2.CandleStore.DefaultCandleIntervals = ""
 
-	sqlConn := &sqlstore.ConnectionSource{
-		Connection: dummyConnection{},
-	}
+	sqlConn := &sqlstore.ConnectionSource{}
+	sqlConn.ToggleTest() // ensure calls to query and copyTo do not fail
 
 	bro, err := broker.New(ctx, logging.NewTestLogger(), conf.Broker, "", eventSource)
 	if err != nil {
@@ -257,18 +253,6 @@ func getTestGRPCServer(t *testing.T, ctx context.Context) (tidy func(), conn *gr
 	return tidy, conn, mockCoreServiceClient, err
 }
 
-type dummyConnection struct {
-	sqlstore.Connection
-}
-
-func (d dummyConnection) Query(context.Context, string, ...interface{}) (pgx.Rows, error) {
-	return nil, pgx.ErrNoRows
-}
-
-func (d dummyConnection) CopyTo(context.Context, io.Writer, string, ...any) (pgconn.CommandTag, error) {
-	return pgconn.CommandTag{}, nil
-}
-
 func TestSubmitTransaction(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()

datanode/sqlstore/connection_tx.go

@@ -16,8 +16,9 @@ import (
 )
 
 type ConnectionSource struct {
-	pool *pgxpool.Pool
-	log  *logging.Logger
+	pool   *pgxpool.Pool
+	log    *logging.Logger
+	isTest bool
 }
 
 type wrappedTx struct {
@@ -36,7 +37,6 @@ type (
 )
 
 func NewTransactionalConnectionSource(ctx context.Context, log *logging.Logger, connConfig ConnectionConfig) (*ConnectionSource, error) {
-	// func NewConnSource(ctx context.Context, log *logging.Logger, connConfig ConnectionConfig) (*ConnectionSource, error) {
 	pool, err := CreateConnectionPool(ctx, connConfig)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create connection pool: %w", err)
@@ -47,6 +47,10 @@ func NewTransactionalConnectionSource(ctx context.Context, log *logging.Logger,
 	}, nil
 }
 
+func (c *ConnectionSource) ToggleTest() {
+	c.isTest = true
+}
+
 func (c *ConnectionSource) WithConnection(ctx context.Context) (context.Context, error) {
 	poolConn, err := c.pool.Acquire(ctx)
 	if err != nil {
@@ -141,6 +145,10 @@ func (c *ConnectionSource) Commit(ctx context.Context) error {
 }
 
 func (c *ConnectionSource) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) {
+	// this is nasty, but required for the API tests currently.
+	if c.isTest && c.pool == nil {
+		return nil, pgx.ErrNoRows
+	}
 	if tx, ok := ctx.Value(txKey{}).(*wrappedTx); ok {
 		return tx.tx.Query(ctx, sql, args...)
 	}
@@ -191,6 +199,10 @@ func (c *ConnectionSource) CopyFrom(ctx context.Context, tableName pgx.Identifie
 }
 
 func (c *ConnectionSource) CopyTo(ctx context.Context, w io.Writer, sql string, args ...any) (pgconn.CommandTag, error) {
+	// this is nasty, but required for the API tests currently.
+	if c.isTest && c.pool == nil {
+		return pgconn.CommandTag{}, nil
+	}
 	var err error
 	sql, err = SanitizeSql(sql, args...)
 	if err != nil {