waf-platform/EdgeNode/internal/utils/kvstore/regression_test.go

package kvstore_test

import (
	"bytes"
	"fmt"
	"github.com/TeaOSLab/EdgeNode/internal/utils/kvstore"
	"strings"
	"testing"
	"time"
)

func openIsolatedTable[T any](t *testing.T, tableName string, encoder kvstore.ValueEncoder[T]) *kvstore.Table[T] {
	storeName := fmt.Sprintf("test-%d", time.Now().UnixNano())

	store, err := kvstore.OpenStore(storeName)
	if err != nil {
		t.Fatal(err)
	}

	t.Cleanup(func() {
		_ = store.Close()
		_ = kvstore.RemoveStore(store.Path())
	})

	db, err := store.NewDB("db1")
	if err != nil {
		t.Fatal(err)
	}

	table, err := kvstore.NewTable[T](tableName, encoder)
	if err != nil {
		t.Fatal(err)
	}

	db.AddTable(table)
	return table
}

func TestQuery_FieldKeyIsStableAcrossPaging(t *testing.T) {
	table := openIsolatedTable[*testCachedItem](t, "cache_items", &testCacheItemEncoder[*testCachedItem]{})

	err := table.AddFields("expiresAt")
	if err != nil {
		t.Fatal(err)
	}

	for i, key := range []string{"a1", "a2", "a3"} {
		err = table.Set(key, &testCachedItem{
			Hash:      key,
			URL:       "https://example.com/" + key,
			ExpiresAt: int64(i + 1),
		})
		if err != nil {
			t.Fatal(err)
		}
	}

	var firstFieldKey []byte
	var firstFieldKeySnapshot []byte
	var count int

	err = table.Query().
		FieldAsc("expiresAt").
		Limit(2).
		FindAll(func(tx *kvstore.Tx[*testCachedItem], item kvstore.Item[*testCachedItem]) (bool, error) {
			switch count {
			case 0:
				firstFieldKey = item.FieldKey
				firstFieldKeySnapshot = append([]byte(nil), item.FieldKey...)
			case 1:
				if !bytes.Equal(firstFieldKey, firstFieldKeySnapshot) {
					t.Fatalf("field key mutated during iteration: got %q want %q", firstFieldKey, firstFieldKeySnapshot)
				}
			}
			count++
			return true, nil
		})
	if err != nil {
		t.Fatal(err)
	}

	var keys []string
	err = table.Query().
		FieldAsc("expiresAt").
		FieldOffset(firstFieldKey).
		Limit(2).
		FindAll(func(tx *kvstore.Tx[*testCachedItem], item kvstore.Item[*testCachedItem]) (bool, error) {
			keys = append(keys, item.Key)
			return true, nil
		})
	if err != nil {
		t.Fatal(err)
	}

	if len(keys) != 2 || keys[0] != "a2" || keys[1] != "a3" {
		t.Fatalf("unexpected paged keys: %v", keys)
	}
}

func TestQuery_FindAll_StringPanicReturnsError(t *testing.T) {
	table := openIsolatedTable[string](t, "users", kvstore.NewStringValueEncoder[string]())

	err := table.Set("a1", "value-1")
	if err != nil {
		t.Fatal(err)
	}

	err = table.Query().
		Limit(1).
		FindAll(func(tx *kvstore.Tx[string], item kvstore.Item[string]) (bool, error) {
			panic("boom")
		})
	if err == nil {
		t.Fatal("expected error, got nil")
	}
	if !strings.Contains(err.Error(), "execute query failed: boom") {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestQuery_ReusesFreshTableTransactionEachRun(t *testing.T) {
	table := openIsolatedTable[string](t, "users", kvstore.NewStringValueEncoder[string]())

	for _, key := range []string{"a1", "a2", "a3"} {
		err := table.Set(key, "value-"+key)
		if err != nil {
			t.Fatal(err)
		}
	}

	query := table.Query().Limit(1)
	for i := 0; i < 2; i++ {
		var keys []string
		err := query.FindAll(func(tx *kvstore.Tx[string], item kvstore.Item[string]) (bool, error) {
			keys = append(keys, item.Key)
			return true, nil
		})
		if err != nil {
			t.Fatal(err)
		}
		if len(keys) != 1 || keys[0] != "a1" {
			t.Fatalf("unexpected result on run %d: %v", i, keys)
		}
	}
}

func TestQuery_UsesExistingTxWithoutClosingIt(t *testing.T) {
	table := openIsolatedTable[string](t, "users", kvstore.NewStringValueEncoder[string]())

	for _, key := range []string{"a1", "a2", "a3"} {
		err := table.Set(key, "value-"+key)
		if err != nil {
			t.Fatal(err)
		}
	}

	err := table.ReadTx(func(tx *kvstore.Tx[string]) error {
		var keys []string
		err := tx.Query().
			Limit(2).
			FindAll(func(queryTx *kvstore.Tx[string], item kvstore.Item[string]) (bool, error) {
				if queryTx != tx {
					return false, fmt.Errorf("query did not reuse the current tx")
				}
				keys = append(keys, item.Key)
				return true, nil
			})
		if err != nil {
			return err
		}
		if len(keys) != 2 {
			return fmt.Errorf("unexpected query size: %d", len(keys))
		}

		_, err = tx.Get("a1")
		return err
	})
	if err != nil {
		t.Fatal(err)
	}
}

func TestDB_InspectProvidesStableBuffers(t *testing.T) {
	table := openIsolatedTable[string](t, "users", kvstore.NewStringValueEncoder[string]())

	for _, pair := range []struct {
		key   string
		value string
	}{
		{key: "a1", value: "value-1"},
		{key: "a2", value: "value-2"},
	} {
		err := table.Set(pair.key, pair.value)
		if err != nil {
			t.Fatal(err)
		}
	}

	var firstKey []byte
	var firstValue []byte
	var firstKeySnapshot []byte
	var firstValueSnapshot []byte
	var count int

	err := table.DB().Inspect(func(key []byte, value []byte) {
		switch count {
		case 0:
			firstKey = key
			firstValue = value
			firstKeySnapshot = append([]byte(nil), key...)
			firstValueSnapshot = append([]byte(nil), value...)
		case 1:
			if !bytes.Equal(firstKey, firstKeySnapshot) {
				t.Fatalf("inspect key mutated after next iteration: got %q want %q", firstKey, firstKeySnapshot)
			}
			if !bytes.Equal(firstValue, firstValueSnapshot) {
				t.Fatalf("inspect value mutated after next iteration: got %q want %q", firstValue, firstValueSnapshot)
			}
		}
		count++
	})
	if err != nil {
		t.Fatal(err)
	}
}