kevo/pkg/iterator/merged_iterator_test.go

package iterator

import (
	"bytes"
	"testing"
)

// mockIterator implements Iterator for testing
type mockIterator struct {
	keys   [][]byte
	values [][]byte
	pos    int
}

func newMockIterator(keys [][]byte, values [][]byte) *mockIterator {
	return &mockIterator{
		keys:   keys,
		values: values,
		pos:    -1, // -1 means not initialized
	}
}

func (m *mockIterator) SeekToFirst() {
	if len(m.keys) > 0 {
		m.pos = 0
	} else {
		m.pos = -1
	}
}

func (m *mockIterator) SeekToLast() {
	if len(m.keys) > 0 {
		m.pos = len(m.keys) - 1
	} else {
		m.pos = -1
	}
}

func (m *mockIterator) Seek(target []byte) bool {
	// Find the first key that is >= target
	for i, key := range m.keys {
		if bytes.Compare(key, target) >= 0 {
			m.pos = i
			return true
		}
	}
	m.pos = -1
	return false
}

func (m *mockIterator) Next() bool {
	if m.pos >= 0 && m.pos < len(m.keys)-1 {
		m.pos++
		return true
	}
	if m.pos == -1 && len(m.keys) > 0 {
		m.pos = 0
		return true
	}
	return false
}

func (m *mockIterator) Key() []byte {
	if m.pos >= 0 && m.pos < len(m.keys) {
		return m.keys[m.pos]
	}
	return nil
}

func (m *mockIterator) Value() []byte {
	if m.pos >= 0 && m.pos < len(m.values) {
		return m.values[m.pos]
	}
	return nil
}

func (m *mockIterator) Valid() bool {
	return m.pos >= 0 && m.pos < len(m.keys)
}

func TestMergedIterator_SeekToFirst(t *testing.T) {
	// Create mock iterators
	iter1 := newMockIterator(
		[][]byte{[]byte("a"), []byte("c"), []byte("e")},
		[][]byte{[]byte("1"), []byte("3"), []byte("5")},
	)
	iter2 := newMockIterator(
		[][]byte{[]byte("b"), []byte("d"), []byte("f")},
		[][]byte{[]byte("2"), []byte("4"), []byte("6")},
	)

	// Create a merged iterator
	merged := NewMergedIterator([]Iterator{iter1, iter2})

	// Test SeekToFirst
	merged.SeekToFirst()
	if !merged.Valid() {
		t.Fatal("Expected iterator to be valid after SeekToFirst")
	}
	if string(merged.Key()) != "a" {
		t.Errorf("Expected first key to be 'a', got '%s'", string(merged.Key()))
	}
	if string(merged.Value()) != "1" {
		t.Errorf("Expected first value to be '1', got '%s'", string(merged.Value()))
	}
}

func TestMergedIterator_Next(t *testing.T) {
	// Create mock iterators
	iter1 := newMockIterator(
		[][]byte{[]byte("a"), []byte("c"), []byte("e")},
		[][]byte{[]byte("1"), []byte("3"), []byte("5")},
	)
	iter2 := newMockIterator(
		[][]byte{[]byte("b"), []byte("d"), []byte("f")},
		[][]byte{[]byte("2"), []byte("4"), []byte("6")},
	)

	// Create a merged iterator
	merged := NewMergedIterator([]Iterator{iter1, iter2})

	// Expected keys and values after merging
	expectedKeys := []string{"a", "b", "c", "d", "e", "f"}
	expectedValues := []string{"1", "2", "3", "4", "5", "6"}

	// Test sequential iteration
	merged.SeekToFirst()
	
	for i, expected := range expectedKeys {
		if !merged.Valid() {
			t.Fatalf("Iterator became invalid at position %d", i)
		}
		if string(merged.Key()) != expected {
			t.Errorf("Expected key at position %d to be '%s', got '%s'", 
				i, expected, string(merged.Key()))
		}
		if string(merged.Value()) != expectedValues[i] {
			t.Errorf("Expected value at position %d to be '%s', got '%s'", 
				i, expectedValues[i], string(merged.Value()))
		}
		if i < len(expectedKeys)-1 && !merged.Next() {
			t.Fatalf("Next() returned false at position %d", i)
		}
	}

	// Test iterating past the end
	if merged.Next() {
		t.Error("Expected Next() to return false after the last key")
	}
}

func TestMergedIterator_Seek(t *testing.T) {
	// Create mock iterators
	iter1 := newMockIterator(
		[][]byte{[]byte("a"), []byte("c"), []byte("e")},
		[][]byte{[]byte("1"), []byte("3"), []byte("5")},
	)
	iter2 := newMockIterator(
		[][]byte{[]byte("b"), []byte("d"), []byte("f")},
		[][]byte{[]byte("2"), []byte("4"), []byte("6")},
	)

	// Create a merged iterator
	merged := NewMergedIterator([]Iterator{iter1, iter2})

	// Test seeking to a position
	if !merged.Seek([]byte("c")) {
		t.Fatal("Expected Seek('c') to return true")
	}
	if string(merged.Key()) != "c" {
		t.Errorf("Expected key after Seek('c') to be 'c', got '%s'", string(merged.Key()))
	}
	if string(merged.Value()) != "3" {
		t.Errorf("Expected value after Seek('c') to be '3', got '%s'", string(merged.Value()))
	}

	// Test seeking to a position that doesn't exist but has a greater key
	if !merged.Seek([]byte("cd")) {
		t.Fatal("Expected Seek('cd') to return true")
	}
	if string(merged.Key()) != "d" {
		t.Errorf("Expected key after Seek('cd') to be 'd', got '%s'", string(merged.Key()))
	}

	// Test seeking beyond the end
	if merged.Seek([]byte("z")) {
		t.Fatal("Expected Seek('z') to return false")
	}
}

func TestMergedIterator_DuplicateKeys(t *testing.T) {
	// Create mock iterators with duplicate keys
	// In a real LSM tree, newer values (from earlier iterators) should take precedence
	iter1 := newMockIterator(
		[][]byte{[]byte("a"), []byte("c")},
		[][]byte{[]byte("newer_a"), []byte("newer_c")},
	)
	iter2 := newMockIterator(
		[][]byte{[]byte("a"), []byte("b")},
		[][]byte{[]byte("older_a"), []byte("b_value")},
	)

	// Create a merged iterator
	merged := NewMergedIterator([]Iterator{iter1, iter2})

	// Test that we get the newer value for key "a"
	merged.SeekToFirst()
	if string(merged.Key()) != "a" {
		t.Errorf("Expected first key to be 'a', got '%s'", string(merged.Key()))
	}
	if string(merged.Value()) != "newer_a" {
		t.Errorf("Expected first value to be 'newer_a', got '%s'", string(merged.Value()))
	}

	// Next should move to "b", skipping the duplicate "a" from iter2
	merged.Next()
	if string(merged.Key()) != "b" {
		t.Errorf("Expected second key to be 'b', got '%s'", string(merged.Key()))
	}
	
	// Then to "c"
	merged.Next()
	if string(merged.Key()) != "c" {
		t.Errorf("Expected third key to be 'c', got '%s'", string(merged.Key()))
	}
}

func TestMergedIterator_SeekToLast(t *testing.T) {
	// Create mock iterators
	iter1 := newMockIterator(
		[][]byte{[]byte("a"), []byte("c"), []byte("e")},
		[][]byte{[]byte("1"), []byte("3"), []byte("5")},
	)
	iter2 := newMockIterator(
		[][]byte{[]byte("b"), []byte("d"), []byte("g")}, // g is the last key
		[][]byte{[]byte("2"), []byte("4"), []byte("7")},
	)

	// Create a merged iterator
	merged := NewMergedIterator([]Iterator{iter1, iter2})

	// Test SeekToLast
	merged.SeekToLast()
	if !merged.Valid() {
		t.Fatal("Expected iterator to be valid after SeekToLast")
	}
	if string(merged.Key()) != "g" {
		t.Errorf("Expected last key to be 'g', got '%s'", string(merged.Key()))
	}
	if string(merged.Value()) != "7" {
		t.Errorf("Expected last value to be '7', got '%s'", string(merged.Value()))
	}
}