//go:build integration
// +build integration

package integration_test

import (
	"database/sql"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/apache/cassandra-gocql-driver/v2"

	"github.com/everything-on-mysql/cassandra-on-mysql/internal/config"
)

func TestStrictProtocolRejectsColumnDefaults(t *testing.T) {
	session := newSession(t, sessionOptions{disableHostLookup: true})

	keyspace := uniqueKeyspace()
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	create := fmt.Sprintf(
		"CREATE TABLE %s.strict_default_reject (id int PRIMARY KEY, created_at timestamp DEFAULT current_timestamp)",
		keyspace,
	)
	err := session.Query(create).Exec()
	if err == nil {
		t.Fatalf("expected strict protocol mode to reject column defaults")
	}
	if !strings.Contains(strings.ToLower(err.Error()), "protocol_mode") {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestExtendedProtocolColumnDefaultsAndOnUpdateDataCorrectness(t *testing.T) {
	requireProxy(t)

	prevCfg := config.GetConfig()
	host, port, stop, err := startProxyWithConfig(mysqlCfg, func(cfg *config.Config) {
		cfg.Executor.ProtocolMode = config.ExecutorProtocolModeExtendedSQL
	})
	if err != nil {
		config.SetConfig(prevCfg)
		t.Fatalf("start proxy with extended protocol failed: %v", err)
	}
	defer func() {
		stop()
		config.SetConfig(prevCfg)
	}()

	session := newSessionForProxy(t, host, port, sessionOptions{disableHostLookup: true})

	keyspace := uniqueKeyspace()
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	table := "extended_defaults"
	createTable := fmt.Sprintf(
		"CREATE TABLE %s.%s (id int PRIMARY KEY, created_at timestamp DEFAULT current_timestamp, updated_at timestamp DEFAULT current_timestamp ON UPDATE current_timestamp, note text)",
		keyspace, table,
	)
	if err := session.Query(createTable).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	insert := fmt.Sprintf("INSERT INTO %s.%s (id, note) VALUES (?, ?)", keyspace, table)
	if err := session.Query(insert, 1, "hello").Exec(); err != nil {
		t.Fatalf("insert failed: %v", err)
	}

	var created1, updated1 time.Time
	var note string
	selectSQL := fmt.Sprintf("SELECT created_at, updated_at, note FROM %s.%s WHERE id = ?", keyspace, table)
	if err := session.Query(selectSQL, 1).Scan(&created1, &updated1, &note); err != nil {
		t.Fatalf("select after insert failed: %v", err)
	}
	if note != "hello" {
		t.Fatalf("unexpected note after insert: %q", note)
	}
	if created1.IsZero() || updated1.IsZero() {
		t.Fatalf("expected default timestamps to be set, got created=%v updated=%v", created1, updated1)
	}

	time.Sleep(20 * time.Millisecond)
	update := fmt.Sprintf("UPDATE %s.%s SET note = ? WHERE id = ?", keyspace, table)
	if err := session.Query(update, "world", 1).Exec(); err != nil {
		t.Fatalf("update failed: %v", err)
	}

	var created2, updated2 time.Time
	if err := session.Query(selectSQL, 1).Scan(&created2, &updated2, &note); err != nil {
		t.Fatalf("select after update failed: %v", err)
	}
	if note != "world" {
		t.Fatalf("unexpected note after update: %q", note)
	}
	if created2.UnixMilli() != created1.UnixMilli() {
		t.Fatalf("created_at should remain unchanged: before=%v after=%v", created1, created2)
	}
	if !updated2.After(updated1) {
		t.Fatalf("updated_at should move forward: before=%v after=%v", updated1, updated2)
	}
}

func TestExtendedProtocolInsertDuplicateIncrementDataCorrectness(t *testing.T) {
	requireProxy(t)

	keyspace := uniqueKeyspace()
	table := "dup_increment"
	targetColumn := fmt.Sprintf("%s.%s.unread_count", keyspace, table)

	prevCfg := config.GetConfig()
	host, port, stop, err := startProxyWithConfig(mysqlCfg, func(cfg *config.Config) {
		cfg.Executor.ProtocolMode = config.ExecutorProtocolModeExtendedSQL
		cfg.Executor.InsertDuplicateIncrementColumns = []string{targetColumn}
	})
	if err != nil {
		config.SetConfig(prevCfg)
		t.Fatalf("start proxy with increment config failed: %v", err)
	}
	defer func() {
		stop()
		config.SetConfig(prevCfg)
	}()

	session := newSessionForProxy(t, host, port, sessionOptions{disableHostLookup: true})
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	createTable := fmt.Sprintf("CREATE TABLE %s.%s (id int PRIMARY KEY, unread_count int, note text)", keyspace, table)
	if err := session.Query(createTable).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	insert := fmt.Sprintf("INSERT INTO %s.%s (id, unread_count, note) VALUES (?, ?, ?)", keyspace, table)
	if err := session.Query(insert, 1, 2, "a").Exec(); err != nil {
		t.Fatalf("first insert failed: %v", err)
	}
	if err := session.Query(insert, 1, 3, "b").Exec(); err != nil {
		t.Fatalf("second insert failed: %v", err)
	}
	if err := session.Query(insert, 1, -2, "c").Exec(); err != nil {
		t.Fatalf("third insert failed: %v", err)
	}

	var unread int
	var note string
	selectSQL := fmt.Sprintf("SELECT unread_count, note FROM %s.%s WHERE id = ?", keyspace, table)
	if err := session.Query(selectSQL, 1).Scan(&unread, &note); err != nil {
		t.Fatalf("select failed: %v", err)
	}
	if unread != 3 {
		t.Fatalf("unexpected unread_count: got=%d want=3", unread)
	}
	if note != "c" {
		t.Fatalf("unexpected note after duplicate inserts: %q", note)
	}
}

func TestExtendedProtocolBatchSamePrimaryKeyInsertDuplicateIncrementDataCorrectness(t *testing.T) {
	requireProxy(t)

	keyspace := uniqueKeyspace()
	table := "batch_dup_increment"
	targetColumn := fmt.Sprintf("%s.%s.unread_count", keyspace, table)

	prevCfg := config.GetConfig()
	host, port, stop, err := startProxyWithConfig(mysqlCfg, func(cfg *config.Config) {
		cfg.Executor.ProtocolMode = config.ExecutorProtocolModeExtendedSQL
		cfg.Executor.InsertDuplicateIncrementColumns = []string{targetColumn}
	})
	if err != nil {
		config.SetConfig(prevCfg)
		t.Fatalf("start proxy with increment config failed: %v", err)
	}
	defer func() {
		stop()
		config.SetConfig(prevCfg)
	}()

	session := newSessionForProxy(t, host, port, sessionOptions{disableHostLookup: true})
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	createTable := fmt.Sprintf("CREATE TABLE %s.%s (id int PRIMARY KEY, unread_count int, note text, content text)", keyspace, table)
	if err := session.Query(createTable).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	seed := fmt.Sprintf("INSERT INTO %s.%s (id, unread_count, note, content) VALUES (?, ?, ?, ?)", keyspace, table)
	if err := session.Query(seed, 1, 10, "seed", "keep").Exec(); err != nil {
		t.Fatalf("seed row failed: %v", err)
	}

	insertUnread := fmt.Sprintf("INSERT INTO %s.%s (id, unread_count) VALUES (?, ?)", keyspace, table)
	insertNote := fmt.Sprintf("INSERT INTO %s.%s (id, note) VALUES (?, ?)", keyspace, table)

	batch := session.NewBatch(gocql.LoggedBatch)
	batch.Query(insertUnread, 1, 2)
	batch.Query(insertUnread, 1, 3)
	batch.Query(insertNote, 1, "batched")
	if err := session.ExecuteBatch(batch); err != nil {
		t.Fatalf("execute batch failed: %v", err)
	}

	var unread int
	var note string
	var content string
	selectSQL := fmt.Sprintf("SELECT unread_count, note, content FROM %s.%s WHERE id = ?", keyspace, table)
	if err := session.Query(selectSQL, 1).Scan(&unread, &note, &content); err != nil {
		t.Fatalf("select after batch failed: %v", err)
	}
	if unread != 15 || note != "batched" || content != "keep" {
		t.Fatalf("unexpected row after duplicate increment batch: unread=%d note=%q content=%q", unread, note, content)
	}
}

func TestExtendedProtocolBatchInsertDuplicateIncrementNullExistingValueDataCorrectness(t *testing.T) {
	requireProxy(t)

	keyspace := uniqueKeyspace()
	table := "batch_dup_increment_null_single"
	targetColumn := fmt.Sprintf("%s.%s.unread_count", keyspace, table)

	prevCfg := config.GetConfig()
	host, port, stop, err := startProxyWithConfig(mysqlCfg, func(cfg *config.Config) {
		cfg.Executor.ProtocolMode = config.ExecutorProtocolModeExtendedSQL
		cfg.Executor.InsertDuplicateIncrementColumns = []string{targetColumn}
	})
	if err != nil {
		config.SetConfig(prevCfg)
		t.Fatalf("start proxy with increment config failed: %v", err)
	}
	defer func() {
		stop()
		config.SetConfig(prevCfg)
	}()

	session := newSessionForProxy(t, host, port, sessionOptions{disableHostLookup: true})
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	createTable := fmt.Sprintf("CREATE TABLE %s.%s (id int PRIMARY KEY, unread_count int, note text)", keyspace, table)
	if err := session.Query(createTable).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	seed := fmt.Sprintf("INSERT INTO %s.%s (id, unread_count, note) VALUES (?, ?, ?)", keyspace, table)
	if err := session.Query(seed, 1, nil, "seed").Exec(); err != nil {
		t.Fatalf("seed row failed: %v", err)
	}

	insert := fmt.Sprintf("INSERT INTO %s.%s (id, unread_count, note) VALUES (?, ?, ?)", keyspace, table)
	batch := session.NewBatch(gocql.LoggedBatch)
	batch.Query(insert, 1, 2, "single")
	if err := session.ExecuteBatch(batch); err != nil {
		t.Fatalf("execute batch failed: %v", err)
	}

	db, err := sql.Open("mysql", mysqlDSN(mysqlCfg))
	if err != nil {
		t.Fatalf("open mysql failed: %v", err)
	}
	defer db.Close()

	var unread sql.NullInt64
	var note string
	query := fmt.Sprintf("SELECT unread_count, note FROM `%s`.`%s` WHERE id = 1", keyspace, table)
	if err := db.QueryRow(query).Scan(&unread, &note); err != nil {
		t.Fatalf("select after batch failed: %v", err)
	}
	if unread.Valid {
		t.Fatalf("expected unread_count to remain NULL, got %d", unread.Int64)
	}
	if note != "single" {
		t.Fatalf("expected note to update to single, got %q", note)
	}
}

func TestExtendedProtocolBatchSamePrimaryKeyInsertDuplicateIncrementNullExistingValueDataCorrectness(t *testing.T) {
	requireProxy(t)

	keyspace := uniqueKeyspace()
	table := "batch_dup_increment_null_same_pk"
	targetColumn := fmt.Sprintf("%s.%s.unread_count", keyspace, table)

	prevCfg := config.GetConfig()
	host, port, stop, err := startProxyWithConfig(mysqlCfg, func(cfg *config.Config) {
		cfg.Executor.ProtocolMode = config.ExecutorProtocolModeExtendedSQL
		cfg.Executor.InsertDuplicateIncrementColumns = []string{targetColumn}
	})
	if err != nil {
		config.SetConfig(prevCfg)
		t.Fatalf("start proxy with increment config failed: %v", err)
	}
	defer func() {
		stop()
		config.SetConfig(prevCfg)
	}()

	session := newSessionForProxy(t, host, port, sessionOptions{disableHostLookup: true})
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	createTable := fmt.Sprintf("CREATE TABLE %s.%s (id int PRIMARY KEY, unread_count int, note text, content text)", keyspace, table)
	if err := session.Query(createTable).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	seed := fmt.Sprintf("INSERT INTO %s.%s (id, unread_count, note, content) VALUES (?, ?, ?, ?)", keyspace, table)
	if err := session.Query(seed, 1, nil, "seed", "keep").Exec(); err != nil {
		t.Fatalf("seed row failed: %v", err)
	}

	insertUnread := fmt.Sprintf("INSERT INTO %s.%s (id, unread_count) VALUES (?, ?)", keyspace, table)
	insertNote := fmt.Sprintf("INSERT INTO %s.%s (id, note) VALUES (?, ?)", keyspace, table)

	batch := session.NewBatch(gocql.LoggedBatch)
	batch.Query(insertUnread, 1, 2)
	batch.Query(insertUnread, 1, 3)
	batch.Query(insertNote, 1, "batched")
	if err := session.ExecuteBatch(batch); err != nil {
		t.Fatalf("execute batch failed: %v", err)
	}

	db, err := sql.Open("mysql", mysqlDSN(mysqlCfg))
	if err != nil {
		t.Fatalf("open mysql failed: %v", err)
	}
	defer db.Close()

	var unread sql.NullInt64
	var note string
	var content string
	query := fmt.Sprintf("SELECT unread_count, note, content FROM `%s`.`%s` WHERE id = 1", keyspace, table)
	if err := db.QueryRow(query).Scan(&unread, &note, &content); err != nil {
		t.Fatalf("select after batch failed: %v", err)
	}
	if unread.Valid {
		t.Fatalf("expected unread_count to remain NULL, got %d", unread.Int64)
	}
	if note != "batched" {
		t.Fatalf("expected note to update to batched, got %q", note)
	}
	if content != "keep" {
		t.Fatalf("expected content to remain keep, got %q", content)
	}
}

func TestExtendedProtocolBatchUpdateMergeDataCorrectness(t *testing.T) {
	requireProxy(t)

	prevCfg := config.GetConfig()
	host, port, stop, err := startProxyWithConfig(mysqlCfg, func(cfg *config.Config) {
		cfg.Executor.ProtocolMode = config.ExecutorProtocolModeExtendedSQL
	})
	if err != nil {
		config.SetConfig(prevCfg)
		t.Fatalf("start proxy with extended protocol failed: %v", err)
	}
	defer func() {
		stop()
		config.SetConfig(prevCfg)
	}()

	session := newSessionForProxy(t, host, port, sessionOptions{disableHostLookup: true})

	keyspace := uniqueKeyspace()
	table := "batch_update_merge"
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	createTable := fmt.Sprintf("CREATE TABLE %s.%s (id int, bucket int, unread_count int, note text, PRIMARY KEY (id, bucket))", keyspace, table)
	if err := session.Query(createTable).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	insert := fmt.Sprintf("INSERT INTO %s.%s (id, bucket, unread_count, note) VALUES (?, ?, ?, ?)", keyspace, table)
	if err := session.Query(insert, 1, 1, 0, "x").Exec(); err != nil {
		t.Fatalf("insert row1 failed: %v", err)
	}
	if err := session.Query(insert, 2, 1, 0, "y").Exec(); err != nil {
		t.Fatalf("insert row2 failed: %v", err)
	}
	if err := session.Query(insert, 3, 1, 0, "z").Exec(); err != nil {
		t.Fatalf("insert row3 failed: %v", err)
	}

	updateWithNote := fmt.Sprintf("UPDATE %s.%s SET unread_count = ?, note = ? WHERE id = ? AND bucket = ?", keyspace, table)
	batch := session.NewBatch(gocql.LoggedBatch)
	batch.Query(updateWithNote, 5, "a", 1, 1)
	batch.Query(updateWithNote, 7, "b", 2, 1)
	if err := session.ExecuteBatch(batch); err != nil {
		t.Fatalf("logged batch update failed: %v", err)
	}

	var unread int
	var note string
	selectRow := fmt.Sprintf("SELECT unread_count, note FROM %s.%s WHERE id = ? AND bucket = ?", keyspace, table)
	if err := session.Query(selectRow, 1, 1).Scan(&unread, &note); err != nil {
		t.Fatalf("select row1 failed: %v", err)
	}
	if unread != 5 || note != "a" {
		t.Fatalf("unexpected row1 after merged batch: unread=%d note=%q", unread, note)
	}
	if err := session.Query(selectRow, 2, 1).Scan(&unread, &note); err != nil {
		t.Fatalf("select row2 failed: %v", err)
	}
	if unread != 7 || note != "b" {
		t.Fatalf("unexpected row2 after merged batch: unread=%d note=%q", unread, note)
	}

	updateWithoutNote := fmt.Sprintf("UPDATE %s.%s SET unread_count = ? WHERE id = ? AND bucket = ?", keyspace, table)
	mixed := session.NewBatch(gocql.LoggedBatch)
	mixed.Query(updateWithNote, 8, "a2", 1, 1)
	mixed.Query(updateWithoutNote, 4, 3, 1)
	if err := session.ExecuteBatch(mixed); err != nil {
		t.Fatalf("mixed logged batch update failed: %v", err)
	}

	if err := session.Query(selectRow, 1, 1).Scan(&unread, &note); err != nil {
		t.Fatalf("select row1 after mixed batch failed: %v", err)
	}
	if unread != 8 || note != "a2" {
		t.Fatalf("unexpected row1 after mixed batch: unread=%d note=%q", unread, note)
	}
	if err := session.Query(selectRow, 3, 1).Scan(&unread, &note); err != nil {
		t.Fatalf("select row3 after mixed batch failed: %v", err)
	}
	if unread != 4 || note != "z" {
		t.Fatalf("unexpected row3 after mixed batch: unread=%d note=%q", unread, note)
	}
}

func TestExtendedProtocolBatchSamePrimaryKeyInsertUpdateDataCorrectness(t *testing.T) {
	requireProxy(t)

	prevCfg := config.GetConfig()
	host, port, stop, err := startProxyWithConfig(mysqlCfg, func(cfg *config.Config) {
		cfg.Executor.ProtocolMode = config.ExecutorProtocolModeExtendedSQL
	})
	if err != nil {
		config.SetConfig(prevCfg)
		t.Fatalf("start proxy with extended protocol failed: %v", err)
	}
	defer func() {
		stop()
		config.SetConfig(prevCfg)
	}()

	session := newSessionForProxy(t, host, port, sessionOptions{disableHostLookup: true})

	keyspace := uniqueKeyspace()
	table := "batch_same_pk_insert_update"
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	createTable := fmt.Sprintf(
		"CREATE TABLE %s.%s (id int, bucket int, unread_count int, note text, content text, PRIMARY KEY (id, bucket))",
		keyspace,
		table,
	)
	if err := session.Query(createTable).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	seed := fmt.Sprintf("INSERT INTO %s.%s (id, bucket, unread_count, note, content) VALUES (?, ?, ?, ?, ?)", keyspace, table)
	if err := session.Query(seed, 1, 1, 5, "seed", "keep").Exec(); err != nil {
		t.Fatalf("seed row failed: %v", err)
	}

	insertNote := fmt.Sprintf("INSERT INTO %s.%s (id, bucket, note) VALUES (?, ?, ?)", keyspace, table)
	updateUnread := fmt.Sprintf("UPDATE %s.%s SET unread_count = unread_count + ? WHERE id = ? AND bucket = ?", keyspace, table)
	updateNote := fmt.Sprintf("UPDATE %s.%s SET note = ? WHERE id = ? AND bucket = ?", keyspace, table)

	batch := session.NewBatch(gocql.LoggedBatch)
	batch.Query(insertNote, 1, 1, "from_insert")
	batch.Query(updateUnread, 2, 1, 1)
	batch.Query(updateNote, "from_update", 1, 1)
	if err := session.ExecuteBatch(batch); err != nil {
		t.Fatalf("execute mixed same-pk batch failed: %v", err)
	}

	var unread int
	var note string
	var content string
	selectSQL := fmt.Sprintf("SELECT unread_count, note, content FROM %s.%s WHERE id = ? AND bucket = ?", keyspace, table)
	if err := session.Query(selectSQL, 1, 1).Scan(&unread, &note, &content); err != nil {
		t.Fatalf("select after batch failed: %v", err)
	}
	if unread != 7 || note != "from_update" || content != "keep" {
		t.Fatalf("unexpected row after same-pk mixed batch: unread=%d note=%q content=%q", unread, note, content)
	}
}

func TestExtendedProtocolBatchArithmeticUpsertDataCorrectness(t *testing.T) {
	requireProxy(t)

	prevCfg := config.GetConfig()
	host, port, stop, err := startProxyWithConfig(mysqlCfg, func(cfg *config.Config) {
		cfg.Executor.ProtocolMode = config.ExecutorProtocolModeExtendedSQL
	})
	if err != nil {
		config.SetConfig(prevCfg)
		t.Fatalf("start proxy with extended protocol failed: %v", err)
	}
	defer func() {
		stop()
		config.SetConfig(prevCfg)
	}()

	session := newSessionForProxy(t, host, port, sessionOptions{disableHostLookup: true})

	keyspace := uniqueKeyspace()
	table := "batch_arithmetic_upsert"
	if err := session.Query(createKeyspaceCQL(keyspace)).Exec(); err != nil {
		t.Fatalf("create keyspace failed: %v", err)
	}
	t.Cleanup(func() {
		_ = session.Query(dropKeyspaceCQL(keyspace)).Exec()
	})

	createTable := fmt.Sprintf(
		"CREATE TABLE %s.%s (user_id int, session_id int, session_type int, unread_count int, mention_count int, note text, PRIMARY KEY (user_id, session_id, session_type))",
		keyspace,
		table,
	)
	if err := session.Query(createTable).Exec(); err != nil {
		t.Fatalf("create table failed: %v", err)
	}

	insertSeed := fmt.Sprintf("INSERT INTO %s.%s (user_id, session_id, session_type, note) VALUES (?, ?, ?, ?)", keyspace, table)
	if err := session.Query(insertSeed, 1, 11, 1, "seed").Exec(); err != nil {
		t.Fatalf("insert seed row failed: %v", err)
	}

	updateArithmetic := fmt.Sprintf(
		"UPDATE %s.%s SET unread_count = unread_count + ?, mention_count = mention_count + ? WHERE user_id = ? AND session_id = ? AND session_type = ?",
		keyspace, table,
	)
	batch := session.NewBatch(gocql.UnloggedBatch)
	batch.Query(updateArithmetic, 3, 2, 1, 11, 1)
	batch.Query(updateArithmetic, 5, 1, 1, 12, 1)
	if err := session.ExecuteBatch(batch); err != nil {
		t.Fatalf("execute arithmetic batch failed: %v", err)
	}

	row1, err := singleRow(session, fmt.Sprintf(
		"SELECT unread_count, mention_count, note FROM %s.%s WHERE user_id = 1 AND session_id = 11 AND session_type = 1",
		keyspace, table,
	))
	if err != nil {
		t.Fatalf("select row1 failed: %v", err)
	}
	unread1, err := parseIntField(row1["unread_count"])
	if err != nil {
		t.Fatalf("parse row1 unread_count failed: %v", err)
	}
	mention1, err := parseIntField(row1["mention_count"])
	if err != nil {
		t.Fatalf("parse row1 mention_count failed: %v", err)
	}
	if unread1 != 3 || mention1 != 2 || asString(row1["note"]) != "seed" {
		t.Fatalf("unexpected row1 after arithmetic batch: %#v", row1)
	}

	row2, err := singleRow(session, fmt.Sprintf(
		"SELECT unread_count, mention_count FROM %s.%s WHERE user_id = 1 AND session_id = 12 AND session_type = 1",
		keyspace, table,
	))
	if err != nil {
		t.Fatalf("select row2 failed: %v", err)
	}
	unread2, err := parseIntField(row2["unread_count"])
	if err != nil {
		t.Fatalf("parse row2 unread_count failed: %v", err)
	}
	mention2, err := parseIntField(row2["mention_count"])
	if err != nil {
		t.Fatalf("parse row2 mention_count failed: %v", err)
	}
	if unread2 != 5 || mention2 != 1 {
		t.Fatalf("unexpected row2 after arithmetic upsert: %#v", row2)
	}

	updateArithmeticOnly := fmt.Sprintf(
		"UPDATE %s.%s SET unread_count = unread_count + ? WHERE user_id = ? AND session_id = ? AND session_type = ?",
		keyspace, table,
	)
	updatePlain := fmt.Sprintf(
		"UPDATE %s.%s SET unread_count = ? WHERE user_id = ? AND session_id = ? AND session_type = ?",
		keyspace, table,
	)
	mixed := session.NewBatch(gocql.UnloggedBatch)
	mixed.Query(updateArithmeticOnly, 1, 1, 11, 1)
	mixed.Query(updatePlain, 100, 1, 12, 1)
	if err := session.ExecuteBatch(mixed); err != nil {
		t.Fatalf("execute mixed arithmetic/plain batch failed: %v", err)
	}

	row1After, err := singleRow(session, fmt.Sprintf(
		"SELECT unread_count FROM %s.%s WHERE user_id = 1 AND session_id = 11 AND session_type = 1",
		keyspace, table,
	))
	if err != nil {
		t.Fatalf("select row1 after mixed batch failed: %v", err)
	}
	unread1After, err := parseIntField(row1After["unread_count"])
	if err != nil {
		t.Fatalf("parse row1 unread_count after mixed batch failed: %v", err)
	}
	if unread1After != 4 {
		t.Fatalf("unexpected row1 unread_count after mixed batch: got=%d want=4", unread1After)
	}

	row2After, err := singleRow(session, fmt.Sprintf(
		"SELECT unread_count FROM %s.%s WHERE user_id = 1 AND session_id = 12 AND session_type = 1",
		keyspace, table,
	))
	if err != nil {
		t.Fatalf("select row2 after mixed batch failed: %v", err)
	}
	unread2After, err := parseIntField(row2After["unread_count"])
	if err != nil {
		t.Fatalf("parse row2 unread_count after mixed batch failed: %v", err)
	}
	if unread2After != 100 {
		t.Fatalf("unexpected row2 unread_count after mixed batch: got=%d want=100", unread2After)
	}
}