kevo/pkg/engine/replication_test.go

package engine

import (
	"os"
	"testing"

	"github.com/jeremytregunna/kevo/pkg/common/clock"
	"github.com/jeremytregunna/kevo/pkg/common/log"
	"github.com/jeremytregunna/kevo/pkg/wal"
)

// TestReplicationHooks tests that the replication hooks are properly called
func TestReplicationHooks(t *testing.T) {
	// Set log level to avoid noise in tests
	log.SetLevel(log.LevelError)
	
	// Create a temporary directory for testing
	tempDir, err := os.MkdirTemp("", "replication-test")
	if err != nil {
		t.Fatalf("Failed to create temporary directory: %v", err)
	}
	defer os.RemoveAll(tempDir)
	
	// Create a new engine
	engine, err := NewEngine(tempDir)
	if err != nil {
		t.Fatalf("Failed to create engine: %v", err)
	}

	// Make sure replication manager was created
	if engine.replicationMgr == nil {
		t.Fatal("Replication manager was not created")
	}
	
	// Verify NodeID was assigned
	if engine.nodeID == (clock.NodeID{}) {
		t.Fatal("NodeID was not assigned")
	}
	
	// Verify Lamport clock was created
	if engine.lamportClock == nil {
		t.Fatal("Lamport clock was not created")
	}
	
	// Test adding and removing replicas
	replicaID := "test-replica"
	fakeNodeID := clock.NodeID{}
	engine.replicationMgr.AddReplica(replicaID, fakeNodeID)
	
	replicas := engine.replicationMgr.GetReplicaIDs()
	if len(replicas) != 1 || replicas[0] != replicaID {
		t.Fatalf("Expected replica ID %s, got %v", replicaID, replicas)
	}
	
	engine.replicationMgr.RemoveReplica(replicaID)
	
	replicas = engine.replicationMgr.GetReplicaIDs()
	if len(replicas) != 0 {
		t.Fatalf("Expected no replicas, got %v", replicas)
	}
	
	// Test setting leader/replica status
	if engine.replicationMgr.IsLeader() {
		t.Fatal("Replication manager should not be leader by default")
	}
	
	engine.replicationMgr.SetLeader(true)
	
	if !engine.replicationMgr.IsLeader() {
		t.Fatal("Replication manager should be leader")
	}
	
	if engine.replicationMgr.IsReplica() {
		t.Fatal("Replication manager should not be replica when it's a leader")
	}
	
	engine.replicationMgr.SetLeader(false)
	
	if !engine.replicationMgr.IsReplica() {
		t.Fatal("Replication manager should be replica when it's not a leader")
	}
	
	// Close the engine
	if err := engine.Close(); err != nil {
		t.Fatalf("Failed to close engine: %v", err)
	}
}

// TestReplicationIntegration tests that the engine works with replication hooks
func TestReplicationIntegration(t *testing.T) {
	// Set log level to avoid noise in tests
	log.SetLevel(log.LevelError)
	
	// Create a temporary directory for testing
	tempDir, err := os.MkdirTemp("", "replication-integration-test")
	if err != nil {
		t.Fatalf("Failed to create temporary directory: %v", err)
	}
	defer os.RemoveAll(tempDir)
	
	// Create a new engine
	engine, err := NewEngine(tempDir)
	if err != nil {
		t.Fatalf("Failed to create engine: %v", err)
	}
	
	// Verify the replication manager is initialized
	if engine.replicationMgr == nil {
		t.Fatal("Replication manager was not initialized")
	}
	
	// Set as leader to ensure replication hooks are called
	engine.replicationMgr.SetLeader(true)
	
	// Perform some operations
	testKey := []byte("test-key")
	testValue := []byte("test-value")
	
	// Put operation
	if err := engine.Put(testKey, testValue); err != nil {
		t.Fatalf("Failed to put: %v", err)
	}
	
	// Get operation
	value, err := engine.Get(testKey)
	if err != nil {
		t.Fatalf("Failed to get: %v", err)
	}
	
	// Verify value
	if string(value) != string(testValue) {
		t.Fatalf("Expected value %q, got %q", testValue, value)
	}
	
	// Close the engine
	if err := engine.Close(); err != nil {
		t.Fatalf("Failed to close engine: %v", err)
	}
}

// TestReplicationInterface verifies that the replication hook interface works correctly
func TestReplicationInterface(t *testing.T) {
	// Set log level to avoid noise in tests
	log.SetLevel(log.LevelError)
	
	// Create a test hook
	callbackCounts := struct {
		singleEntries int
		batchEntries  int
	}{}
	
	testHook := &syncTestReplicationHook{
		onEntryCallback: func(entry *wal.Entry, ts clock.Timestamp) error {
			callbackCounts.singleEntries++
			return nil
		},
		onBatchCallback: func(entries []*wal.Entry, ts clock.Timestamp) error {
			callbackCounts.batchEntries += len(entries)
			return nil
		},
	}
	
	// Create a test entry
	entry := &wal.Entry{
		SequenceNumber: 1,
		Type:           wal.OpTypePut,
		Key:            []byte("key1"),
		Value:          []byte("value1"),
	}
	
	// Create a test batch
	batch := []*wal.Entry{
		{
			SequenceNumber: 2,
			Type:           wal.OpTypePut,
			Key:            []byte("key2"),
			Value:          []byte("value2"),
		},
		{
			SequenceNumber: 3,
			Type:           wal.OpTypePut,
			Key:            []byte("key3"),
			Value:          []byte("value3"),
		},
	}
	
	// Create a timestamp
	nodeID := clock.NodeID{}
	ts := clock.Timestamp{
		Counter: 1,
		Node:    nodeID,
	}
	
	// Call the hook methods directly
	if err := testHook.OnEntryWritten(entry, ts); err != nil {
		t.Fatalf("OnEntryWritten failed: %v", err)
	}
	
	if err := testHook.OnBatchWritten(batch, ts); err != nil {
		t.Fatalf("OnBatchWritten failed: %v", err)
	}
	
	// Verify callbacks were called
	if callbackCounts.singleEntries != 1 {
		t.Errorf("Expected 1 single entry callback, got %d", callbackCounts.singleEntries)
	}
	
	if callbackCounts.batchEntries != 2 {
		t.Errorf("Expected 2 batch entry callbacks, got %d", callbackCounts.batchEntries)
	}
}

// Test helper: a synchronous mock replication hook for testing
type syncTestReplicationHook struct {
	onEntryCallback func(*wal.Entry, clock.Timestamp) error
	onBatchCallback func([]*wal.Entry, clock.Timestamp) error
}

func (h *syncTestReplicationHook) OnEntryWritten(entry *wal.Entry, ts clock.Timestamp) error {
	if h.onEntryCallback != nil {
		return h.onEntryCallback(entry, ts)
	}
	return nil
}

func (h *syncTestReplicationHook) OnBatchWritten(entries []*wal.Entry, ts clock.Timestamp) error {
	if h.onBatchCallback != nil {
		return h.onBatchCallback(entries, ts)
	}
	return nil
}

// Test helper: an asynchronous mock replication hook for testing
type testReplicationHook struct {
	onEntryCallback func(*wal.Entry, clock.Timestamp) error
	onBatchCallback func([]*wal.Entry, clock.Timestamp) error
}

func (h *testReplicationHook) OnEntryWritten(entry *wal.Entry, ts clock.Timestamp) error {
	if h.onEntryCallback != nil {
		return h.onEntryCallback(entry, ts)
	}
	return nil
}

func (h *testReplicationHook) OnBatchWritten(entries []*wal.Entry, ts clock.Timestamp) error {
	if h.onBatchCallback != nil {
		return h.onBatchCallback(entries, ts)
	}
	return nil
}