blockchain/internal/lib/cassandra/table_test.go

package cassandra

import (
	"context"
	"testing"

	"github.com/scylladb/gocqlx/v3/qb"
	"github.com/stretchr/testify/assert"
)

func TestQueryBuilder_TableDriven(t *testing.T) {
	ctx := context.Background()
	db := &DB{} // 可以用 mock DB

	type args struct {
		cmp      qb.Cmp
		whereArg map[string]any
		selects  []string
		orderCol string
		order    qb.Order
		limit    uint
		setCol   string
		setVal   any
	}

	tests := []struct {
		name         string
		args         args
		wantPanic    bool
		wantColumns  []string
		wantOrderCol string
		wantOrder    qb.Order
		wantLimit    uint
		wantSetCol   string
		wantSetVal   any
	}{
		{
			name: "where by partition key",
			args: args{
				cmp:      qb.Eq("id"),
				whereArg: map[string]any{"id": "abc"},
				selects:  []string{"id", "name"},
				orderCol: "id",
				order:    qb.ASC,
				limit:    1,
				setCol:   "name",
				setVal:   "Daniel",
			},
			wantPanic:    false,
			wantColumns:  []string{"id", "name"},
			wantOrderCol: "id",
			wantOrder:    qb.ASC,
			wantLimit:    1,
			wantSetCol:   "name",
			wantSetVal:   "Daniel",
		},
		{
			name: "where by sai index",
			args: args{
				cmp:      qb.Eq("name"),
				whereArg: map[string]any{"name": "daniel"},
				selects:  []string{"id", "name"},
				orderCol: "name",
				order:    qb.DESC,
				limit:    2,
				setCol:   "name",
				setVal:   "Jacky",
			},
			wantPanic:    false,
			wantColumns:  []string{"id", "name"},
			wantOrderCol: "name",
			wantOrder:    qb.DESC,
			wantLimit:    2,
			wantSetCol:   "name",
			wantSetVal:   "Jacky",
		},
		{
			name: "where by non-partition-non-sai",
			args: args{
				cmp:      qb.Eq("age"),
				whereArg: map[string]any{"age": 18},
				selects:  []string{"id", "name"},
				orderCol: "age",
				order:    qb.ASC,
				limit:    3,
				setCol:   "age",
				setVal:   20,
			},
			wantPanic: true,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			run := func() {
				q := db.Model(ctx, &MonkeyEntity{}, "my_keyspace").
					Where(tc.args.cmp, tc.args.whereArg).
					Select(tc.args.selects...).
					OrderBy(tc.args.orderCol, tc.args.order).
					Limit(tc.args.limit).
					Set(tc.args.setCol, tc.args.setVal)

				assert.Equal(t, tc.wantColumns, q.columns)
				if len(q.orders) > 0 {
					assert.Equal(t, tc.wantOrderCol, q.orders[0].Column)
					assert.Equal(t, tc.wantOrder, q.orders[0].Order)
				}
				assert.Equal(t, tc.wantLimit, q.limit)
				if len(q.sets) > 0 {
					assert.Equal(t, tc.wantSetCol, q.sets[0].Col)
					assert.Equal(t, tc.wantSetVal, q.sets[0].Val)
				}
			}

			if tc.wantPanic {
				assert.Panics(t, run)
			} else {
				assert.NotPanics(t, run)
			}
		})
	}
}

func TestQuery_Select_TableDriven(t *testing.T) {
	tests := []struct {
		name        string
		selectCalls [][]string
		wantColumns []string
	}{
		{
			name:        "select one col",
			selectCalls: [][]string{{"id"}},
			wantColumns: []string{"id"},
		},
		{
			name:        "select multi col in one call",
			selectCalls: [][]string{{"id", "name"}},
			wantColumns: []string{"id", "name"},
		},
		{
			name:        "multiple select calls append columns",
			selectCalls: [][]string{{"id"}, {"name"}, {"age"}},
			wantColumns: []string{"id", "name", "age"},
		},
		{
			name:        "multiple select calls with overlap",
			selectCalls: [][]string{{"id"}, {"id", "name"}, {"name", "age"}},
			wantColumns: []string{"id", "id", "name", "name", "age"},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			q := &Query{columns: make([]string, 0)}
			for _, call := range tc.selectCalls {
				q = q.Select(call...)
			}
			assert.Equal(t, tc.wantColumns, q.columns)
		})
	}
}