AetherForge/server/internal/api/fleet_torrent_test.go

package api

import (
	"encoding/json"
	"testing"
	"time"

	"crypto-miner-server/internal/atlas"
	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"

	"github.com/gorilla/websocket"
)

func TestFleetTorrentGossipRelayCrossSubnet(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	hub.SetServerPolicy(ServerPolicy{FleetTorrentEnabled: true})
	_ = database.UpsertAgent(&models.Agent{ID: "ft-a", Name: "a", IP: "10.1.1.10", Status: "online"})
	_ = database.UpsertAgent(&models.Agent{ID: "ft-b", Name: "b", IP: "10.2.2.20", Status: "online"})

	connA := connectTestAgentWithIP(t, hub, "ft-a", "10.1.1.10")
	connB := connectTestAgentWithIP(t, hub, "ft-b", "10.2.2.20")

	recvCh := make(chan Message, 2)
	go readUntilType(connB, "fleet_torrent_gossip", recvCh)

	payload, _ := json.Marshal(map[string]interface{}{
		"records": []atlas.FleetGossipRecord{{
			Kind: atlas.FleetGossipHaveShard, AgentID: "ft-a", Token: "tok",
			ShardIndex: 0, ShardHash: "abc", Subnet: "10.1.1",
		}},
	})
	if err := connA.WriteJSON(Message{Type: "fleet_torrent_gossip", Payload: payload}); err != nil {
		t.Fatal(err)
	}

	select {
	case msg := <-recvCh:
		var body map[string]interface{}
		if err := json.Unmarshal(msg.Payload, &body); err != nil {
			t.Fatal(err)
		}
		recs, _ := body["records"].([]interface{})
		if len(recs) != 1 {
			t.Fatalf("records=%v", body["records"])
		}
	case <-time.After(2 * time.Second):
		t.Fatal("cross-subnet peer did not receive fleet_torrent_gossip")
	}
	_ = websocket.CloseNormalClosure
}

func TestAuthSubnetPrimarySeederHint(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	hub.SetServerPolicy(ServerPolicy{FleetRolesEnabled: true, FleetTorrentEnabled: true})

	conn, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn, map[string]interface{}{
		"agent_id":    "seed-primary-aa",
		"hostname":    "host",
		"platform":    "windows",
		"version":     "test",
		"fleet_role":  "seeder",
		"seeder_mode": true,
	})
	var body map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &body); err != nil {
		t.Fatal(err)
	}
	if _, ok := body["subnet_primary_seeder"].(string); !ok {
		t.Fatalf("subnet_primary_seeder missing: %#v", body)
	}
	if body["fleet_torrent_enabled"] != true {
		t.Fatalf("fleet_torrent_enabled=%#v", body["fleet_torrent_enabled"])
	}
}

func TestFleetTorrentGossipSameVPCScope(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	hub.SetServerPolicy(ServerPolicy{FleetTorrentEnabled: true})
	hub.storeAgentCloudMeta("vpc-a", CloudInstanceMeta{VpcID: "vpc-42", Region: "us-east-1"})
	hub.storeAgentCloudMeta("vpc-b", CloudInstanceMeta{VpcID: "vpc-42", Region: "us-east-1"})
	_ = database.UpsertAgent(&models.Agent{ID: "vpc-a", Name: "a", IP: "10.10.1.1", Status: "online"})
	_ = database.UpsertAgent(&models.Agent{ID: "vpc-b", Name: "b", IP: "10.20.2.2", Status: "online"})

	if scope := primarySeederScope("10.10.1.1", hub.agentCloudMetaLocked("vpc-a")); scope != "vpc-42" {
		t.Fatalf("scope=%q", scope)
	}
	if !hub.agentMatchesPrimaryScopeLocked("vpc-a", "vpc-42") {
		t.Fatal("expected vpc-a in vpc-42")
	}
	if hub.agentMatchesPrimaryScopeLocked("vpc-b", "vpc-99") {
		t.Fatal("expected vpc-b mismatch on other vpc")
	}
}

func TestSubnetPrimarySeederElection(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	hub.storeAgentFleetRole("aaa-seeder", "seeder")
	hub.storeAgentFleetRole("bbb-seeder", "seeder")
	hub.mu.Lock()
	hub.agentLiveTelemetry["aaa-seeder"] = map[string]interface{}{"fleet_role": "seeder"}
	hub.agentLiveTelemetry["bbb-seeder"] = map[string]interface{}{"fleet_role": "seeder"}
	hub.mu.Unlock()
	_ = database.UpsertAgent(&models.Agent{ID: "aaa-seeder", IP: "10.5.5.1"})
	_ = database.UpsertAgent(&models.Agent{ID: "bbb-seeder", IP: "10.5.5.2"})
	pick := hub.electSubnetPrimarySeeder("10.5.5")
	if pick != "aaa-seeder" {
		t.Fatalf("pick=%q", pick)
	}
}