atomic_set_test.go

package bmdb

import (
	"bytes"
	"encoding/binary"
	"testing"
)

// decodeSnapshotHashes decodes the optimized binary format produced by
// AtomicSet.SnapshotHashes: [N:uint32] then N * { klen:uint16, key, m:uint32, m * uint32(hash) }
func decodeSnapshotHashes(b []byte) (map[string][]uint32, error) {
	res := make(map[string][]uint32)
	r := bytes.NewReader(b)
	var n uint32
	if err := binary.Read(r, binary.BigEndian, &n); err != nil {
		return nil, err
	}
	for i := uint32(0); i < n; i++ {
		var klen uint16
		if err := binary.Read(r, binary.BigEndian, &klen); err != nil {
			return nil, err
		}
		kb := make([]byte, klen)
		if _, err := r.Read(kb); err != nil {
			return nil, err
		}
		var m uint32
		if err := binary.Read(r, binary.BigEndian, &m); err != nil {
			return nil, err
		}
		arr := make([]uint32, 0, m)
		for j := uint32(0); j < m; j++ {
			var h uint32
			if err := binary.Read(r, binary.BigEndian, &h); err != nil {
				return nil, err
			}
			arr = append(arr, h)
		}
		res[string(kb)] = arr
	}
	return res, nil
}

func TestAtomicSet_SnapshotAndReplace(t *testing.T) {
	s := NewSet()
	// create some members
	v1 := []byte("a")
	v2 := []byte("b")
	r1 := &Record{ // minimal record; fields not relevant for hashing test
		Key:   []byte("k1"),
		Value: v1,
	}
	r2 := &Record{
		Key:   []byte("k1"),
		Value: v2,
	}
	// add to concrete set
	if err := s.SAdd("key1", [][]byte{v1, v2}, []*Record{r1, r2}); err != nil {
		t.Fatalf("SAdd failed: %v", err)
	}

	a := NewAtomicSet(s)

	// snapshot should encode the key1 and two hashes
	b, err := a.SnapshotHashes()
	if err != nil {
		t.Fatalf("SnapshotHashes failed: %v", err)
	}

	decoded, err := decodeSnapshotHashes(b)
	if err != nil {
		t.Fatalf("decodeSnapshotHashes failed: %v", err)
	}

	if arr, ok := decoded["key1"]; !ok {
		t.Fatalf("expected key1 in snapshot")
	} else if len(arr) != 2 {
		t.Fatalf("expected 2 members for key1, got %d", len(arr))
	}

	// Now modify concrete set and ReplaceFromSet
	// remove one member
	if err := s.SRem("key1", v2); err != nil {
		t.Fatalf("SRem failed: %v", err)
	}
	a.ReplaceFromSet(s)

	b2, err := a.SnapshotHashes()
	if err != nil {
		t.Fatalf("SnapshotHashes after replace failed: %v", err)
	}
	decoded2, err := decodeSnapshotHashes(b2)
	if err != nil {
		t.Fatalf("decodeSnapshotHashes 2 failed: %v", err)
	}
	if arr, ok := decoded2["key1"]; !ok {
		t.Fatalf("expected key1 in snapshot after replace")
	} else if len(arr) != 1 {
		t.Fatalf("expected 1 member for key1 after replace, got %d", len(arr))
	}
}

func TestAtomicSetBasic(t *testing.T) {
	s := NewSet()
	rec := &Record{Key: []byte("k"), Value: []byte("v")}
	if err := s.SAdd("myset", [][]byte{[]byte("v")}, []*Record{rec}); err != nil {
		t.Fatalf("SAdd failed: %v", err)
	}

	as := NewAtomicSet(s)
	// membership should be visible
	ok, err := as.SIsMember("myset", []byte("v"))
	if err != nil || !ok {
		t.Fatalf("expected member exist")
	}

	// add via atomic wrapper and check
	rec2 := &Record{Key: []byte("k2"), Value: []byte("v2")}
	if err := as.SAdd("myset", [][]byte{[]byte("v2")}, []*Record{rec2}); err != nil {
		t.Fatalf("SAdd via atomic failed: %v", err)
	}
	ok2, _ := as.SIsMember("myset", []byte("v2"))
	if !ok2 {
		t.Fatalf("expected v2 present after SAdd")
	}
}