CHORUS/pkg/dht/dht_test.go

package dht

import (
	"context"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p"
	"github.com/libp2p/go-libp2p/core/host"
	"github.com/libp2p/go-libp2p/core/test"
	dht "github.com/libp2p/go-libp2p-kad-dht"
	"github.com/multiformats/go-multiaddr"
)

func TestDefaultConfig(t *testing.T) {
	config := DefaultConfig()
	
	if config.ProtocolPrefix != "/CHORUS" {
		t.Errorf("expected protocol prefix '/CHORUS', got %s", config.ProtocolPrefix)
	}
	
	if config.BootstrapTimeout != 30*time.Second {
		t.Errorf("expected bootstrap timeout 30s, got %v", config.BootstrapTimeout)
	}
	
	if config.Mode != dht.ModeAuto {
		t.Errorf("expected mode auto, got %v", config.Mode)
	}
	
	if !config.AutoBootstrap {
		t.Error("expected auto bootstrap to be enabled")
	}
}

func TestNewDHT(t *testing.T) {
	ctx := context.Background()
	
	// Create a test host
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	// Test with default options
	d, err := NewDHT(ctx, host)
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	if d.host != host {
		t.Error("host not set correctly")
	}
	
	if d.config.ProtocolPrefix != "/CHORUS" {
		t.Errorf("expected protocol prefix '/CHORUS', got %s", d.config.ProtocolPrefix)
	}
}

func TestDHTWithOptions(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	// Test with custom options
	d, err := NewDHT(ctx, host,
		WithProtocolPrefix("/custom"),
		WithMode(dht.ModeClient),
		WithBootstrapTimeout(60*time.Second),
		WithDiscoveryInterval(120*time.Second),
		WithAutoBootstrap(false),
	)
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	if d.config.ProtocolPrefix != "/custom" {
		t.Errorf("expected protocol prefix '/custom', got %s", d.config.ProtocolPrefix)
	}
	
	if d.config.Mode != dht.ModeClient {
		t.Errorf("expected mode client, got %v", d.config.Mode)
	}
	
	if d.config.BootstrapTimeout != 60*time.Second {
		t.Errorf("expected bootstrap timeout 60s, got %v", d.config.BootstrapTimeout)
	}
	
	if d.config.DiscoveryInterval != 120*time.Second {
		t.Errorf("expected discovery interval 120s, got %v", d.config.DiscoveryInterval)
	}
	
	if d.config.AutoBootstrap {
		t.Error("expected auto bootstrap to be disabled")
	}
}

func TestWithBootstrapPeersFromStrings(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	bootstrapAddrs := []string{
		"/ip4/127.0.0.1/tcp/4001/p2p/QmTest1",
		"/ip4/127.0.0.1/tcp/4002/p2p/QmTest2",
	}
	
	d, err := NewDHT(ctx, host, WithBootstrapPeersFromStrings(bootstrapAddrs))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	if len(d.config.BootstrapPeers) != 2 {
		t.Errorf("expected 2 bootstrap peers, got %d", len(d.config.BootstrapPeers))
	}
}

func TestWithBootstrapPeersFromStringsInvalid(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	// Include invalid addresses - they should be filtered out
	bootstrapAddrs := []string{
		"/ip4/127.0.0.1/tcp/4001/p2p/QmTest1", // valid
		"invalid-address",                      // invalid
		"/ip4/127.0.0.1/tcp/4002/p2p/QmTest2", // valid
	}
	
	d, err := NewDHT(ctx, host, WithBootstrapPeersFromStrings(bootstrapAddrs))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Should have filtered out the invalid address
	if len(d.config.BootstrapPeers) != 2 {
		t.Errorf("expected 2 valid bootstrap peers, got %d", len(d.config.BootstrapPeers))
	}
}

func TestBootstrapWithoutPeers(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host, WithAutoBootstrap(false))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Bootstrap should use default IPFS peers when none configured
	err = d.Bootstrap()
	// This might fail in test environment without network access, but should not panic
	if err != nil {
		// Expected in test environment
		t.Logf("Bootstrap failed as expected in test environment: %v", err)
	}
}

func TestIsBootstrapped(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host, WithAutoBootstrap(false))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Should not be bootstrapped initially
	if d.IsBootstrapped() {
		t.Error("DHT should not be bootstrapped initially")
	}
}

func TestRegisterPeer(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host)
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	peerID := test.RandPeerIDFatal(t)
	agent := "claude"
	role := "frontend"
	capabilities := []string{"react", "javascript"}
	
	d.RegisterPeer(peerID, agent, role, capabilities)
	
	knownPeers := d.GetKnownPeers()
	if len(knownPeers) != 1 {
		t.Errorf("expected 1 known peer, got %d", len(knownPeers))
	}
	
	peerInfo, exists := knownPeers[peerID]
	if !exists {
		t.Error("peer not found in known peers")
	}
	
	if peerInfo.Agent != agent {
		t.Errorf("expected agent %s, got %s", agent, peerInfo.Agent)
	}
	
	if peerInfo.Role != role {
		t.Errorf("expected role %s, got %s", role, peerInfo.Role)
	}
	
	if len(peerInfo.Capabilities) != len(capabilities) {
		t.Errorf("expected %d capabilities, got %d", len(capabilities), len(peerInfo.Capabilities))
	}
}

func TestGetConnectedPeers(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host)
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Initially should have no connected peers
	peers := d.GetConnectedPeers()
	if len(peers) != 0 {
		t.Errorf("expected 0 connected peers, got %d", len(peers))
	}
}

func TestPutAndGetValue(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host, WithAutoBootstrap(false))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Test without bootstrap (should fail)
	key := "test-key"
	value := []byte("test-value")
	
	err = d.PutValue(ctx, key, value)
	if err == nil {
		t.Error("PutValue should fail when DHT not bootstrapped")
	}
	
	_, err = d.GetValue(ctx, key)
	if err == nil {
		t.Error("GetValue should fail when DHT not bootstrapped")
	}
}

func TestProvideAndFindProviders(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host, WithAutoBootstrap(false))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Test without bootstrap (should fail)
	key := "test-service"
	
	err = d.Provide(ctx, key)
	if err == nil {
		t.Error("Provide should fail when DHT not bootstrapped")
	}
	
	_, err = d.FindProviders(ctx, key, 10)
	if err == nil {
		t.Error("FindProviders should fail when DHT not bootstrapped")
	}
}

func TestFindPeer(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host, WithAutoBootstrap(false))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Test without bootstrap (should fail)
	peerID := test.RandPeerIDFatal(t)
	
	_, err = d.FindPeer(ctx, peerID)
	if err == nil {
		t.Error("FindPeer should fail when DHT not bootstrapped")
	}
}

func TestFindPeersByRole(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host, WithAutoBootstrap(false))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Register some local peers
	peerID1 := test.RandPeerIDFatal(t)
	peerID2 := test.RandPeerIDFatal(t)
	
	d.RegisterPeer(peerID1, "claude", "frontend", []string{"react"})
	d.RegisterPeer(peerID2, "claude", "backend", []string{"go"})
	
	// Find frontend peers
	frontendPeers, err := d.FindPeersByRole(ctx, "frontend")
	if err != nil {
		t.Fatalf("failed to find peers by role: %v", err)
	}
	
	if len(frontendPeers) != 1 {
		t.Errorf("expected 1 frontend peer, got %d", len(frontendPeers))
	}
	
	if frontendPeers[0].ID != peerID1 {
		t.Error("wrong peer returned for frontend role")
	}
	
	// Find all peers with wildcard
	allPeers, err := d.FindPeersByRole(ctx, "*")
	if err != nil {
		t.Fatalf("failed to find all peers: %v", err)
	}
	
	if len(allPeers) != 2 {
		t.Errorf("expected 2 peers with wildcard, got %d", len(allPeers))
	}
}

func TestAnnounceRole(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host, WithAutoBootstrap(false))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Should fail when not bootstrapped
	err = d.AnnounceRole(ctx, "frontend")
	if err == nil {
		t.Error("AnnounceRole should fail when DHT not bootstrapped")
	}
}

func TestAnnounceCapability(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host, WithAutoBootstrap(false))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Should fail when not bootstrapped
	err = d.AnnounceCapability(ctx, "react")
	if err == nil {
		t.Error("AnnounceCapability should fail when DHT not bootstrapped")
	}
}

func TestGetRoutingTable(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host)
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	rt := d.GetRoutingTable()
	if rt == nil {
		t.Error("routing table should not be nil")
	}
}

func TestGetDHTSize(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host)
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	size := d.GetDHTSize()
	// Should be 0 or small initially
	if size < 0 {
		t.Errorf("DHT size should be non-negative, got %d", size)
	}
}

func TestRefreshRoutingTable(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host, WithAutoBootstrap(false))
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	// Should fail when not bootstrapped
	err = d.RefreshRoutingTable()
	if err == nil {
		t.Error("RefreshRoutingTable should fail when DHT not bootstrapped")
	}
}

func TestHost(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host)
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	defer d.Close()
	
	if d.Host() != host {
		t.Error("Host() should return the same host instance")
	}
}

func TestClose(t *testing.T) {
	ctx := context.Background()
	
	host, err := libp2p.New()
	if err != nil {
		t.Fatalf("failed to create test host: %v", err)
	}
	defer host.Close()
	
	d, err := NewDHT(ctx, host)
	if err != nil {
		t.Fatalf("failed to create DHT: %v", err)
	}
	
	// Should close without error
	err = d.Close()
	if err != nil {
		t.Errorf("Close() failed: %v", err)
	}
}