CHORUS/pkg/protocol/resolver_test.go

package protocol

import (
	"context"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p/core/peer"
	"github.com/libp2p/go-libp2p/core/peerstore"
	"github.com/libp2p/go-libp2p/core/test"
)

func TestNewResolver(t *testing.T) {
	// Create a mock peerstore
	mockPeerstore := &mockPeerstore{}
	
	resolver := NewResolver(mockPeerstore)
	
	if resolver == nil {
		t.Fatal("resolver is nil")
	}
	
	if resolver.peerstore != mockPeerstore {
		t.Error("peerstore not set correctly")
	}
	
	if resolver.defaultStrategy != StrategyBestMatch {
		t.Errorf("expected default strategy %v, got %v", StrategyBestMatch, resolver.defaultStrategy)
	}
	
	if resolver.maxPeersPerResult != 5 {
		t.Errorf("expected max peers per result 5, got %d", resolver.maxPeersPerResult)
	}
}

func TestResolverWithOptions(t *testing.T) {
	mockPeerstore := &mockPeerstore{}
	
	resolver := NewResolver(mockPeerstore,
		WithCacheTTL(10*time.Minute),
		WithDefaultStrategy(StrategyPriority),
		WithMaxPeersPerResult(10),
	)
	
	if resolver.cacheTTL != 10*time.Minute {
		t.Errorf("expected cache TTL 10m, got %v", resolver.cacheTTL)
	}
	
	if resolver.defaultStrategy != StrategyPriority {
		t.Errorf("expected strategy %v, got %v", StrategyPriority, resolver.defaultStrategy)
	}
	
	if resolver.maxPeersPerResult != 10 {
		t.Errorf("expected max peers 10, got %d", resolver.maxPeersPerResult)
	}
}

func TestRegisterPeer(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{})
	
	peerID := test.RandPeerIDFatal(t)
	capability := &PeerCapability{
		Agent:          "claude",
		Role:           "frontend",
		Capabilities:   []string{"react", "javascript"},
		Models:         []string{"claude-3"},
		Specialization: "frontend",
		Status:         "ready",
		Metadata:       make(map[string]string),
	}
	
	resolver.RegisterPeer(peerID, capability)
	
	// Verify peer was registered
	caps := resolver.GetPeerCapabilities()
	if len(caps) != 1 {
		t.Errorf("expected 1 peer, got %d", len(caps))
	}
	
	registeredCap, exists := caps[peerID]
	if !exists {
		t.Error("peer not found in capabilities")
	}
	
	if registeredCap.Agent != capability.Agent {
		t.Errorf("expected agent %s, got %s", capability.Agent, registeredCap.Agent)
	}
	
	if registeredCap.PeerID != peerID {
		t.Error("peer ID not set correctly")
	}
}

func TestUnregisterPeer(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{})
	
	peerID := test.RandPeerIDFatal(t)
	capability := &PeerCapability{
		Agent: "claude",
		Role:  "frontend",
	}
	
	// Register then unregister
	resolver.RegisterPeer(peerID, capability)
	resolver.UnregisterPeer(peerID)
	
	caps := resolver.GetPeerCapabilities()
	if len(caps) != 0 {
		t.Errorf("expected 0 peers after unregister, got %d", len(caps))
	}
}

func TestUpdatePeerStatus(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{})
	
	peerID := test.RandPeerIDFatal(t)
	capability := &PeerCapability{
		Agent:  "claude",
		Role:   "frontend",
		Status: "ready",
	}
	
	resolver.RegisterPeer(peerID, capability)
	resolver.UpdatePeerStatus(peerID, "busy")
	
	caps := resolver.GetPeerCapabilities()
	updatedCap := caps[peerID]
	
	if updatedCap.Status != "busy" {
		t.Errorf("expected status 'busy', got '%s'", updatedCap.Status)
	}
}

func TestResolveURI(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{})
	
	// Register some test peers
	peerID1 := test.RandPeerIDFatal(t)
	peerID2 := test.RandPeerIDFatal(t)
	
	resolver.RegisterPeer(peerID1, &PeerCapability{
		Agent:        "claude",
		Role:         "frontend",
		Capabilities: []string{"react", "javascript"},
		Status:       "ready",
		Metadata:     map[string]string{"project": "chorus"},
	})
	
	resolver.RegisterPeer(peerID2, &PeerCapability{
		Agent:        "claude",
		Role:         "backend",
		Capabilities: []string{"go", "api"},
		Status:       "ready",
		Metadata:     map[string]string{"project": "chorus"},
	})
	
	// Test exact match
	uri, err := ParseBzzzURI("bzzz://claude:frontend@chorus:react")
	if err != nil {
		t.Fatalf("failed to parse URI: %v", err)
	}
	
	ctx := context.Background()
	result, err := resolver.Resolve(ctx, uri)
	if err != nil {
		t.Fatalf("failed to resolve URI: %v", err)
	}
	
	if len(result.Peers) != 1 {
		t.Errorf("expected 1 peer in result, got %d", len(result.Peers))
	}
	
	if result.Peers[0].PeerID != peerID1 {
		t.Error("wrong peer returned")
	}
}

func TestResolveURIWithWildcards(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{})
	
	peerID1 := test.RandPeerIDFatal(t)
	peerID2 := test.RandPeerIDFatal(t)
	
	resolver.RegisterPeer(peerID1, &PeerCapability{
		Agent:        "claude",
		Role:         "frontend",
		Capabilities: []string{"react"},
		Status:       "ready",
	})
	
	resolver.RegisterPeer(peerID2, &PeerCapability{
		Agent:        "claude", 
		Role:         "backend",
		Capabilities: []string{"go"},
		Status:       "ready",
	})
	
	// Test wildcard match
	uri, err := ParseBzzzURI("bzzz://claude:*@*:*")
	if err != nil {
		t.Fatalf("failed to parse URI: %v", err)
	}
	
	ctx := context.Background()
	result, err := resolver.Resolve(ctx, uri)
	if err != nil {
		t.Fatalf("failed to resolve URI: %v", err)
	}
	
	if len(result.Peers) != 2 {
		t.Errorf("expected 2 peers in result, got %d", len(result.Peers))
	}
}

func TestResolveURIWithOfflinePeers(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{})
	
	peerID := test.RandPeerIDFatal(t)
	
	resolver.RegisterPeer(peerID, &PeerCapability{
		Agent:  "claude",
		Role:   "frontend",
		Status: "offline", // This peer should be filtered out
	})
	
	uri, err := ParseBzzzURI("bzzz://claude:frontend@*:*")
	if err != nil {
		t.Fatalf("failed to parse URI: %v", err)
	}
	
	ctx := context.Background()
	result, err := resolver.Resolve(ctx, uri)
	if err != nil {
		t.Fatalf("failed to resolve URI: %v", err)
	}
	
	if len(result.Peers) != 0 {
		t.Errorf("expected 0 peers (offline filtered), got %d", len(result.Peers))
	}
}

func TestResolveString(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{})
	
	peerID := test.RandPeerIDFatal(t)
	resolver.RegisterPeer(peerID, &PeerCapability{
		Agent:  "claude",
		Role:   "frontend",
		Status: "ready",
	})
	
	ctx := context.Background()
	result, err := resolver.ResolveString(ctx, "bzzz://claude:frontend@*:*")
	if err != nil {
		t.Fatalf("failed to resolve string: %v", err)
	}
	
	if len(result.Peers) != 1 {
		t.Errorf("expected 1 peer, got %d", len(result.Peers))
	}
}

func TestResolverCaching(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{}, WithCacheTTL(1*time.Second))
	
	peerID := test.RandPeerIDFatal(t)
	resolver.RegisterPeer(peerID, &PeerCapability{
		Agent:  "claude",
		Role:   "frontend",
		Status: "ready",
	})
	
	ctx := context.Background()
	uri := "bzzz://claude:frontend@*:*"
	
	// First resolution should hit the resolver
	result1, err := resolver.ResolveString(ctx, uri)
	if err != nil {
		t.Fatalf("failed to resolve: %v", err)
	}
	
	// Second resolution should hit the cache
	result2, err := resolver.ResolveString(ctx, uri)
	if err != nil {
		t.Fatalf("failed to resolve: %v", err)
	}
	
	// Results should be identical (from cache)
	if result1.ResolvedAt != result2.ResolvedAt {
		// This is expected behavior - cache should return same timestamp
	}
	
	// Wait for cache to expire
	time.Sleep(2 * time.Second)
	
	// Third resolution should miss cache and create new result
	result3, err := resolver.ResolveString(ctx, uri)
	if err != nil {
		t.Fatalf("failed to resolve: %v", err)
	}
	
	if result3.ResolvedAt.Before(result1.ResolvedAt.Add(1 * time.Second)) {
		t.Error("cache should have expired and created new result")
	}
}

func TestResolutionStrategies(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{})
	
	// Register peers with different priorities
	peerID1 := test.RandPeerIDFatal(t)
	peerID2 := test.RandPeerIDFatal(t)
	
	resolver.RegisterPeer(peerID1, &PeerCapability{
		Agent:  "claude",
		Role:   "frontend",
		Status: "ready",
	})
	
	resolver.RegisterPeer(peerID2, &PeerCapability{
		Agent:  "claude",
		Role:   "frontend", 
		Status: "busy",
	})
	
	ctx := context.Background()
	uri, _ := ParseBzzzURI("bzzz://claude:frontend@*:*")
	
	// Test different strategies
	strategies := []ResolutionStrategy{
		StrategyBestMatch,
		StrategyPriority,
		StrategyLoadBalance,
		StrategyExact,
	}
	
	for _, strategy := range strategies {
		result, err := resolver.Resolve(ctx, uri, strategy)
		if err != nil {
			t.Errorf("failed to resolve with strategy %s: %v", strategy, err)
		}
		
		if len(result.Peers) == 0 {
			t.Errorf("no peers found with strategy %s", strategy)
		}
		
		if result.Strategy != string(strategy) {
			t.Errorf("strategy not recorded correctly: expected %s, got %s", strategy, result.Strategy)
		}
	}
}

func TestPeerMatching(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{})
	
	capability := &PeerCapability{
		Agent:        "claude",
		Role:         "frontend",
		Capabilities: []string{"react", "javascript"},
		Status:       "ready",
		Metadata:     map[string]string{"project": "chorus"},
	}
	
	tests := []struct {
		name     string
		uri      *BzzzURI
		expected bool
	}{
		{
			name:     "exact match",
			uri:      &BzzzURI{Agent: "claude", Role: "frontend", Project: "chorus", Task: "react"},
			expected: true,
		},
		{
			name:     "wildcard agent",
			uri:      &BzzzURI{Agent: "*", Role: "frontend", Project: "chorus", Task: "react"},
			expected: true,
		},
		{
			name:     "capability match",
			uri:      &BzzzURI{Agent: "claude", Role: "frontend", Project: "*", Task: "javascript"},
			expected: true,
		},
		{
			name:     "no match - wrong agent",
			uri:      &BzzzURI{Agent: "gpt", Role: "frontend", Project: "chorus", Task: "react"},
			expected: false,
		},
		{
			name:     "no match - wrong role",
			uri:      &BzzzURI{Agent: "claude", Role: "backend", Project: "chorus", Task: "react"},
			expected: false,
		},
	}
	
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := resolver.peerMatches(capability, tt.uri)
			if result != tt.expected {
				t.Errorf("expected %v, got %v", tt.expected, result)
			}
		})
	}
}

func TestGetPeerCapability(t *testing.T) {
	resolver := NewResolver(&mockPeerstore{})
	
	peerID := test.RandPeerIDFatal(t)
	capability := &PeerCapability{
		Agent: "claude",
		Role:  "frontend",
	}
	
	// Test before registration
	_, exists := resolver.GetPeerCapability(peerID)
	if exists {
		t.Error("peer should not exist before registration")
	}
	
	// Register and test
	resolver.RegisterPeer(peerID, capability)
	
	retrieved, exists := resolver.GetPeerCapability(peerID)
	if !exists {
		t.Error("peer should exist after registration")
	}
	
	if retrieved.Agent != capability.Agent {
		t.Errorf("expected agent %s, got %s", capability.Agent, retrieved.Agent)
	}
}

// Mock peerstore implementation for testing
type mockPeerstore struct{}

func (m *mockPeerstore) PeerInfo(peer.ID) peer.AddrInfo                    { return peer.AddrInfo{} }
func (m *mockPeerstore) Peers() peer.IDSlice                               { return nil }
func (m *mockPeerstore) Addrs(peer.ID) []peerstore.Multiaddr              { return nil }
func (m *mockPeerstore) AddrStream(context.Context, peer.ID) <-chan peerstore.Multiaddr { return nil }
func (m *mockPeerstore) SetAddr(peer.ID, peerstore.Multiaddr, time.Duration) {}
func (m *mockPeerstore) SetAddrs(peer.ID, []peerstore.Multiaddr, time.Duration) {}
func (m *mockPeerstore) UpdateAddrs(peer.ID, time.Duration, time.Duration) {}
func (m *mockPeerstore) ClearAddrs(peer.ID) {}
func (m *mockPeerstore) PeersWithAddrs() peer.IDSlice { return nil }
func (m *mockPeerstore) PubKey(peer.ID) peerstore.PubKey { return nil }
func (m *mockPeerstore) SetPubKey(peer.ID, peerstore.PubKey) error { return nil }
func (m *mockPeerstore) PrivKey(peer.ID) peerstore.PrivKey { return nil }
func (m *mockPeerstore) SetPrivKey(peer.ID, peerstore.PrivKey) error { return nil }
func (m *mockPeerstore) Get(peer.ID, string) (interface{}, error) { return nil, nil }
func (m *mockPeerstore) Put(peer.ID, string, interface{}) error { return nil }
func (m *mockPeerstore) GetProtocols(peer.ID) ([]peerstore.Protocol, error) { return nil, nil }
func (m *mockPeerstore) SetProtocols(peer.ID, ...peerstore.Protocol) error { return nil }
func (m *mockPeerstore) SupportsProtocols(peer.ID, ...peerstore.Protocol) ([]peerstore.Protocol, error) { return nil, nil }
func (m *mockPeerstore) RemovePeer(peer.ID) {}
func (m *mockPeerstore) Close() error { return nil }