go/maintnotifications/e2e/scenario_unified_injector_test.go

package e2e

import (
	"context"
	"fmt"
	"os"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/redis/go-redis/v9"
	"github.com/redis/go-redis/v9/maintnotifications"
)

// TestUnifiedInjector_SMIGRATING demonstrates using the unified notification injector
// This test requires the proxy-based mock to directly inject notifications
func TestUnifiedInjector_SMIGRATING(t *testing.T) {
	// Skip if using real fault injector (these tests require proxy mock for direct notification injection)
	if os.Getenv("FAULT_INJECTOR_URL") != "" {
		t.Skip("Skipping unified injector test - requires proxy mock, not real fault injector")
	}

	ctx := context.Background()

	// Create notification injector (automatically chooses based on environment)
	injector, err := NewNotificationInjector()
	if err != nil {
		t.Fatalf("[ERROR] Failed to create notification injector: %v", err)
	}

	// Start the injector - skip if proxy not available
	if err := injector.Start(); err != nil {
		t.Skipf("Skipping test - proxy not available: %v", err)
	}
	defer injector.Stop()

	// Get test mode configuration
	testMode := injector.GetTestModeConfig()
	t.Logf("Using %s mode", testMode.Mode)
	t.Logf("Cluster addresses: %v", injector.GetClusterAddrs())

	// Create cluster client with maintnotifications enabled
	client := redis.NewClusterClient(&redis.ClusterOptions{
		Addrs:    injector.GetClusterAddrs(),
		Protocol: 3, // RESP3 required for push notifications

		MaintNotificationsConfig: &maintnotifications.Config{
			Mode:           maintnotifications.ModeEnabled,
			RelaxedTimeout: 30 * time.Second,
		},
	})
	defer client.Close()

	// Set up notification tracking
	tracker := NewTrackingNotificationsHook()
	setupNotificationHook(client, tracker)

	// Verify connection
	if err := client.Ping(ctx).Err(); err != nil {
		t.Fatalf("[ERROR] Failed to connect to cluster: %v", err)
	}

	// Perform some operations to establish connections
	for i := 0; i < 10; i++ {
		if err := client.Set(ctx, fmt.Sprintf("key%d", i), "value", 0).Err(); err != nil {
			t.Logf("Warning: Failed to set key: %v", err)
		}
	}

	// Start a blocking operation in background to keep connection active
	// This ensures the proxy has an active connection to send notifications to
	blockingDone := make(chan error, 1)
	go func() {
		// BLPOP with 10 second timeout - keeps connection active
		_, err := client.BLPop(ctx, 10*time.Second, "notification-test-list").Result()
		blockingDone <- err
	}()

	// Wait for blocking command to start (mode-aware)
	time.Sleep(testMode.ConnectionEstablishDelay)

	// Inject SMIGRATING notification while connection is active
	t.Log("Injecting SMIGRATING notification...")
	if err := injector.InjectSMIGRATING(ctx, 12345, "1000-2000", "3000"); err != nil {
		t.Fatalf("[ERROR] Failed to inject SMIGRATING: %v", err)
	}

	// Wait for notification processing (mode-aware)
	time.Sleep(testMode.NotificationDelay)

	// Verify notification was received
	analysis := tracker.GetAnalysis()
	if analysis.SMigratingCount == 0 {
		t.Errorf("[ERROR] Expected to receive SMIGRATING notification, got 0")
	} else {
		t.Logf("✓ Received %d SMIGRATING notification(s)", analysis.SMigratingCount)
	}

	// Verify operations still work (timeouts should be relaxed)
	if err := client.Set(ctx, "test-key-during-migration", "value", 0).Err(); err != nil {
		t.Errorf("[ERROR] Expected operations to work during migration, got error: %v", err)
	}

	// Print analysis
	analysis.Print(t)

	t.Log("✓ SMIGRATING test passed")
}

// TestUnifiedInjector_SMIGRATED demonstrates SMIGRATED notification handling
// This test requires the proxy-based mock to directly inject notifications
func TestUnifiedInjector_SMIGRATED(t *testing.T) {
	// Skip if using real fault injector (these tests require proxy mock for direct notification injection)
	if os.Getenv("FAULT_INJECTOR_URL") != "" {
		t.Skip("Skipping unified injector test - requires proxy mock, not real fault injector")
	}

	ctx := context.Background()

	injector, err := NewNotificationInjector()
	if err != nil {
		t.Fatalf("[ERROR] Failed to create notification injector: %v", err)
	}

	// Start the injector - skip if proxy not available
	if err := injector.Start(); err != nil {
		t.Skipf("Skipping test - proxy not available: %v", err)
	}
	defer injector.Stop()

	// Get test mode configuration
	testMode := injector.GetTestModeConfig()
	t.Logf("Using %s mode", testMode.Mode)

	// Track cluster state reloads
	var reloadCount atomic.Int32

	client := redis.NewClusterClient(&redis.ClusterOptions{
		Addrs:    injector.GetClusterAddrs(),
		Protocol: 3,

		MaintNotificationsConfig: &maintnotifications.Config{
			Mode:           maintnotifications.ModeEnabled,
			RelaxedTimeout: 30 * time.Second,
		},
	})
	defer client.Close()

	// Set up notification tracking
	tracker := NewTrackingNotificationsHook()
	setupNotificationHook(client, tracker)

	if err := client.Ping(ctx).Err(); err != nil {
		t.Fatalf("[ERROR] Failed to connect: %v", err)
	}

	// Set up reload callback on existing nodes
	client.ForEachShard(ctx, func(ctx context.Context, nodeClient *redis.Client) error {
		manager := nodeClient.GetMaintNotificationsManager()
		if manager != nil {
			manager.SetClusterStateReloadCallback(func(ctx context.Context, hostPort string, slotRanges []string) {
				reloadCount.Add(1)
				t.Logf("Cluster state reload triggered for %s, slots: %v", hostPort, slotRanges)
			})
		}
		return nil
	})

	// Set up reload callback for new nodes
	client.OnNewNode(func(nodeClient *redis.Client) {
		manager := nodeClient.GetMaintNotificationsManager()
		if manager != nil {
			manager.SetClusterStateReloadCallback(func(ctx context.Context, hostPort string, slotRanges []string) {
				reloadCount.Add(1)
				t.Logf("Cluster state reload triggered for %s, slots: %v", hostPort, slotRanges)
			})
		}
	})

	// Perform some operations to establish connections
	for i := 0; i < 10; i++ {
		if err := client.Set(ctx, fmt.Sprintf("key%d", i), "value", 0).Err(); err != nil {
			t.Logf("Warning: Failed to set key: %v", err)
		}
	}

	initialReloads := reloadCount.Load()

	// Start a blocking operation in background to keep connection active
	blockingDone := make(chan error, 1)
	go func() {
		// BLPOP with 10 second timeout - keeps connection active
		_, err := client.BLPop(ctx, 10*time.Second, "notification-test-list").Result()
		blockingDone <- err
	}()

	// Wait for blocking command to start (mode-aware)
	time.Sleep(testMode.ConnectionEstablishDelay)

	// Get all node addresses - needed for both modes
	addrs := injector.GetClusterAddrs()
	var newNodeAddr string
	if len(addrs) >= 4 {
		newNodeAddr = addrs[3] // Node 3: 127.0.0.1:17003
	} else {
		// Fallback to first node if we don't have 4 nodes
		newNodeAddr = addrs[0]
	}

	// Mode-specific behavior for triggering SMIGRATED
	if testMode.IsProxyMock() {
		// Proxy mock: Directly inject SMIGRATED notification
		t.Log("Injecting SMIGRATED notification to swap node 2 for node 3...")
		t.Logf("Using new node address: %s", newNodeAddr)

		if err := injector.InjectSMIGRATED(ctx, 12346, newNodeAddr, "1000-2000", "3000"); err != nil {
			t.Fatalf("[ERROR] Failed to inject SMIGRATED: %v", err)
		}
	} else {
		// Real fault injector: Trigger slot migration which generates both SMIGRATING and SMIGRATED
		t.Log("Triggering slot migration (will generate SMIGRATING and SMIGRATED)...")

		// First inject SMIGRATING (this triggers the actual migration)
		if err := injector.InjectSMIGRATING(ctx, 12345, "1000-2000", "3000"); err != nil {
			t.Fatalf("[ERROR] Failed to trigger slot migration: %v", err)
		}

		// Wait for migration to complete (real FI takes longer)
		t.Log("Waiting for migration to complete...")
	}

	// Wait for notification processing (mode-aware)
	time.Sleep(testMode.NotificationDelay)

	// Wait for blocking operation to complete
	<-blockingDone

	// Verify notification was received
	analysis := tracker.GetAnalysis()

	if testMode.IsProxyMock() {
		// Proxy mock: SMIGRATED notifications may not always be received
		// because they're sent to all connections, but the client might not be actively
		// listening on all of them. This is expected behavior.
		if analysis.MigratedCount > 0 {
			t.Logf("✓ Received %d SMIGRATED notification(s)", analysis.MigratedCount)
		} else {
			t.Logf("Note: No SMIGRATED notifications received (expected in proxy mock mode)")
		}
	} else {
		// Real FI: Should receive both SMIGRATING and SMIGRATED
		if analysis.MigratingCount == 0 {
			t.Errorf("[ERROR] Expected to receive SMIGRATING notification with real FI, got 0")
		} else {
			t.Logf("✓ Received %d SMIGRATING notification(s)", analysis.MigratingCount)
		}

		if analysis.MigratedCount > 0 {
			t.Logf("✓ Received %d SMIGRATED notification(s)", analysis.MigratedCount)
		} else {
			t.Logf("Note: SMIGRATED notification not yet received (migration may still be in progress)")
		}
	}

	// Verify cluster state reload callback
	// Note: SMIGRATED notifications trigger cluster reload via the endpoint information
	// However, in proxy mock mode, the callback might not be triggered because
	// the notification is sent to all connections, not just the active one
	finalReloads := reloadCount.Load()
	if finalReloads > initialReloads {
		t.Logf("✓ Cluster state reloaded %d time(s)", finalReloads-initialReloads)
	} else {
		t.Logf("Note: Cluster state reload callback not triggered (expected in proxy mock mode)")
	}

	// Verify the client discovered the new node topology
	// After SMIGRATED, the client should reload cluster slots
	// Give it a moment to process
	time.Sleep(2 * time.Second)

	// Count how many nodes the client knows about
	var nodeAddrsMu sync.Mutex
	nodeAddrs := make(map[string]bool)
	client.ForEachShard(ctx, func(ctx context.Context, nodeClient *redis.Client) error {
		addr := nodeClient.Options().Addr
		nodeAddrsMu.Lock()
		nodeAddrs[addr] = true
		nodeAddrsMu.Unlock()
		t.Logf("Client knows about node: %s", addr)
		return nil
	})

	// We should have 3 nodes (0, 1, 3) after the swap
	if len(nodeAddrs) < 3 {
		t.Logf("Warning: Expected client to discover 3 nodes after SMIGRATED, got %d", len(nodeAddrs))
	} else {
		t.Logf("✓ Client discovered %d node(s) after SMIGRATED", len(nodeAddrs))
	}

	// Verify the new node (17003) is in the list
	if len(addrs) >= 4 && !nodeAddrs[newNodeAddr] {
		t.Logf("Warning: Client did not discover new node %s", newNodeAddr)
	} else if len(addrs) >= 4 {
		t.Logf("✓ Client discovered new node %s", newNodeAddr)
	}

	// Verify we can still perform operations after SMIGRATED
	// The client should have reloaded cluster topology
	successCount := 0
	for i := 0; i < 10; i++ {
		key := fmt.Sprintf("post-migration-key-%d", i)
		if err := client.Set(ctx, key, "value", 0).Err(); err == nil {
			successCount++
		} else {
			t.Logf("Warning: Failed to set key after SMIGRATED: %v", err)
		}
	}

	if successCount < 8 {
		t.Errorf("[ERROR] Expected most operations to succeed after SMIGRATED, got %d/10", successCount)
	} else {
		t.Logf("✓ Successfully performed %d/10 operations after SMIGRATED", successCount)
	}

	// Print analysis
	analysis.Print(t)

	t.Log("✓ SMIGRATED test passed")
}

// TestUnifiedInjector_ComplexScenario demonstrates a complex migration scenario
// This test requires the proxy-based mock to directly inject notifications
func TestUnifiedInjector_ComplexScenario(t *testing.T) {
	// Skip if using real fault injector (these tests require proxy mock for direct notification injection)
	if os.Getenv("FAULT_INJECTOR_URL") != "" {
		t.Skip("Skipping unified injector test - requires proxy mock, not real fault injector")
	}

	ctx := context.Background()

	injector, err := NewNotificationInjector()
	if err != nil {
		t.Fatalf("[ERROR] Failed to create notification injector: %v", err)
	}

	// Start the injector - skip if proxy not available
	if err := injector.Start(); err != nil {
		t.Skipf("Skipping test - proxy not available: %v", err)
	}
	defer injector.Stop()

	// Get test mode configuration
	testMode := injector.GetTestModeConfig()
	t.Logf("Using %s mode", testMode.Mode)

	var reloadCount atomic.Int32

	client := redis.NewClusterClient(&redis.ClusterOptions{
		Addrs:    injector.GetClusterAddrs(),
		Protocol: 3,
	})
	defer client.Close()

	tracker := NewTrackingNotificationsHook()
	setupNotificationHook(client, tracker)

	client.OnNewNode(func(nodeClient *redis.Client) {
		manager := nodeClient.GetMaintNotificationsManager()
		if manager != nil {
			manager.SetClusterStateReloadCallback(func(ctx context.Context, hostPort string, slotRanges []string) {
				reloadCount.Add(1)
			})
		}
	})

	if err := client.Ping(ctx).Err(); err != nil {
		t.Fatalf("[ERROR] Failed to connect: %v", err)
	}

	// Perform operations
	for i := 0; i < 20; i++ {
		client.Set(ctx, fmt.Sprintf("key%d", i), "value", 0)
	}

	// Wait for connections to establish (mode-aware)
	time.Sleep(testMode.ConnectionEstablishDelay)

	// Simulate a multi-step migration scenario
	t.Log("Step 1: Injecting SMIGRATING for slots 0-5000...")
	if err := injector.InjectSMIGRATING(ctx, 10001, "0-5000"); err != nil {
		t.Fatalf("[ERROR] Failed to inject SMIGRATING: %v", err)
	}

	// Wait for notification processing (mode-aware)
	time.Sleep(testMode.NotificationDelay)

	// Verify operations work during migration
	for i := 0; i < 5; i++ {
		if err := client.Set(ctx, fmt.Sprintf("migration-key%d", i), "value", 0).Err(); err != nil {
			t.Logf("Warning: Operation failed during migration: %v", err)
		}
	}

	if testMode.IsProxyMock() {
		// Only inject SMIGRATED with mock injector
		t.Log("Step 2: Injecting SMIGRATED for completed migration...")
		addrs := injector.GetClusterAddrs()
		hostPort := addrs[0]

		if err := injector.InjectSMIGRATED(ctx, 10002, hostPort, "0-5000"); err != nil {
			t.Fatalf("[ERROR] Failed to inject SMIGRATED: %v", err)
		}

		// Wait for notification processing (mode-aware)
		time.Sleep(testMode.NotificationDelay)
	}

	// Verify operations still work
	for i := 0; i < 5; i++ {
		if err := client.Set(ctx, fmt.Sprintf("post-migration-key%d", i), "value", 0).Err(); err != nil {
			t.Errorf("[ERROR] Operations failed after migration: %v", err)
		}
	}

	// Print final analysis
	analysis := tracker.GetAnalysis()
	analysis.Print(t)

	t.Logf("✓ Complex scenario test passed")
	t.Logf("  - SMIGRATING notifications: %d", analysis.MigratingCount)
	t.Logf("  - SMIGRATED notifications: %d", analysis.MigratedCount)
	t.Logf("  - Cluster state reloads: %d", reloadCount.Load())
}