From d7a99dbfc5636fa87a59a83cdde0c0b6f683c6ad Mon Sep 17 00:00:00 2001
From: Ian Suvak <ian.suvak@avalabs.org>
Date: Wed, 16 Oct 2024 14:33:17 -0400
Subject: [PATCH] add test for connected weight calculation

---
 peers/app_request_network.go      | 33 ++++++++++------
 peers/app_request_network_test.go | 66 +++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+), 11 deletions(-)
 create mode 100644 peers/app_request_network_test.go

diff --git a/peers/app_request_network.go b/peers/app_request_network.go
index 76ad46f6..7b2d1978 100644
--- a/peers/app_request_network.go
+++ b/peers/app_request_network.go
@@ -250,17 +250,7 @@ func (n *appRequestNetwork) ConnectToCanonicalValidators(subnetID ids.ID) (*Conn
 	connectedNodes := n.ConnectPeers(nodeIDs)
 
 	// Calculate the total weight of connected validators.
-	connectedBLSPubKeys := set.NewSet[string](len(validatorSet))
-	connectedWeight := uint64(0)
-	for node := range connectedNodes {
-		vdr := validatorSet[nodeValidatorIndexMap[node]]
-		blsPubKey := hex.EncodeToString(vdr.PublicKeyBytes)
-		if connectedBLSPubKeys.Contains(blsPubKey) {
-			continue
-		}
-		connectedWeight += vdr.Weight
-		connectedBLSPubKeys.Add(blsPubKey)
-	}
+	connectedWeight := calculateConnectedWeight(validatorSet, nodeValidatorIndexMap, connectedNodes)
 
 	return &ConnectedCanonicalValidators{
 		ConnectedWeight:       connectedWeight,
@@ -300,3 +290,24 @@ func (n *appRequestNetwork) setInfoAPICallLatencyMS(latency float64) {
 func (n *appRequestNetwork) setPChainAPICallLatencyMS(latency float64) {
 	n.metrics.pChainAPICallLatencyMS.Observe(latency)
 }
+
+// Non-receiver util functions
+
+func calculateConnectedWeight(
+	validatorSet []*warp.Validator,
+	nodeValidatorIndexMap map[ids.NodeID]int,
+	connectedNodes set.Set[ids.NodeID],
+) uint64 {
+	connectedBLSPubKeys := set.NewSet[string](len(validatorSet))
+	connectedWeight := uint64(0)
+	for node := range connectedNodes {
+		vdr := validatorSet[nodeValidatorIndexMap[node]]
+		blsPubKey := hex.EncodeToString(vdr.PublicKeyBytes)
+		if connectedBLSPubKeys.Contains(blsPubKey) {
+			continue
+		}
+		connectedWeight += vdr.Weight
+		connectedBLSPubKeys.Add(blsPubKey)
+	}
+	return connectedWeight
+}
diff --git a/peers/app_request_network_test.go b/peers/app_request_network_test.go
new file mode 100644
index 00000000..acd690d4
--- /dev/null
+++ b/peers/app_request_network_test.go
@@ -0,0 +1,66 @@
+// Copyright (C) 2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package peers
+
+import (
+	"testing"
+
+	"github.com/ava-labs/avalanchego/ids"
+	"github.com/ava-labs/avalanchego/utils/crypto/bls"
+	"github.com/ava-labs/avalanchego/utils/set"
+	"github.com/ava-labs/avalanchego/vms/platformvm/warp"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCalculateConnectedWeight(t *testing.T) {
+	vdr1 := makeValidator(t, 10, 1)
+	vdr2 := makeValidator(t, 20, 1)
+	vdr3 := makeValidator(t, 30, 2)
+	vdrs := []*warp.Validator{&vdr1, &vdr2, &vdr3}
+	nodeValidatorIndexMap := map[ids.NodeID]int{
+		vdr1.NodeIDs[0]: 0,
+		vdr2.NodeIDs[0]: 1,
+		vdr3.NodeIDs[0]: 2,
+		vdr3.NodeIDs[1]: 2,
+	}
+	var connectedNodes set.Set[ids.NodeID]
+	connectedNodes.Add(vdr1.NodeIDs[0])
+	connectedNodes.Add(vdr2.NodeIDs[0])
+
+	// vdr1 and vdr2 are connected, so their weight should be added
+	require.Equal(t, 2, connectedNodes.Len())
+	connectedWeight := calculateConnectedWeight(vdrs, nodeValidatorIndexMap, connectedNodes)
+	require.Equal(t, uint64(30), connectedWeight)
+
+	// Add one of the vdr3's nodeIDs to the connected nodes
+	// and confirm that it adds vdr3's weight
+	connectedNodes.Add(vdr3.NodeIDs[0])
+	require.Equal(t, 3, connectedNodes.Len())
+	connectedWeight2 := calculateConnectedWeight(vdrs, nodeValidatorIndexMap, connectedNodes)
+	require.Equal(t, uint64(60), connectedWeight2)
+
+	// Add another of vdr3's nodeIDs to the connected nodes
+	// and confirm that it's weight isn't double counted
+	connectedNodes.Add(vdr3.NodeIDs[1])
+	require.Equal(t, 4, connectedNodes.Len())
+	connectedWeight3 := calculateConnectedWeight(vdrs, nodeValidatorIndexMap, connectedNodes)
+	require.Equal(t, uint64(60), connectedWeight3)
+}
+
+func makeValidator(t *testing.T, weight uint64, numNodeIDs int) warp.Validator {
+	sk, err := bls.NewSecretKey()
+	require.NoError(t, err)
+	pk := bls.PublicFromSecretKey(sk)
+
+	nodeIDs := make([]ids.NodeID, numNodeIDs)
+	for i := 0; i < numNodeIDs; i++ {
+		nodeIDs[i] = ids.GenerateTestNodeID()
+	}
+	return warp.Validator{
+		PublicKey:      pk,
+		PublicKeyBytes: bls.PublicKeyToUncompressedBytes(pk),
+		Weight:         weight,
+		NodeIDs:        nodeIDs,
+	}
+}