rpctest: Add NodesConnected

This allows package users to check whether two nodes are connected.

rpctest/rpc_harness_test.go

@@ -202,7 +202,75 @@ func testDisconnectNode(r *Harness, t *testing.T) {
 	// Ensure the nodes remain connected after trying to disconnect them in the
 	// reverse order.
 	assertConnectedTo(t, harness, r)
+}
 
+func testNodesConnected(r *Harness, t *testing.T) {
+	// Create a fresh test harness.
+	harness, err := New(&chaincfg.RegNetParams, nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := harness.SetUp(false, 0); err != nil {
+		t.Fatalf("unable to complete rpctest setup: %v", err)
+	}
+	defer harness.TearDown()
+
+	// Establish a p2p connection from our new local harness to the main
+	// harness.
+	if err := ConnectNode(harness, r); err != nil {
+		t.Fatalf("unable to connect local to main harness: %v", err)
+	}
+
+	// Sanity check.
+	assertConnectedTo(t, harness, r)
+
+	// Ensure nodes are still connected.
+	assertConnectedTo(t, harness, r)
+
+	testCases := []struct {
+		name         string
+		allowReverse bool
+		expected     bool
+		from         *Harness
+		to           *Harness
+	}{
+		// The existing connection is h->r.
+		{"!allowReverse, h->r", false, true, harness, r},
+		{"allowReverse, h->r", true, true, harness, r},
+		{"!allowReverse, r->h", false, false, r, harness},
+		{"allowReverse, r->h", true, true, r, harness},
+	}
+
+	for _, tc := range testCases {
+		actual, err := NodesConnected(tc.from, tc.to, tc.allowReverse)
+		if err != nil {
+			t.Fatalf("unable to determine node connection: %v", err)
+		}
+		if actual != tc.expected {
+			t.Fatalf("test case %s: actual result (%v) differs from expected "+
+				"(%v)", tc.name, actual, tc.expected)
+		}
+	}
+
+	// Disconnect the nodes.
+	if err := RemoveNode(harness, r); err != nil {
+		t.Fatalf("unable to disconnect local to main harness: %v", err)
+	}
+
+	// Sanity check.
+	assertNotConnectedTo(t, harness, r)
+
+	// All test cases must return false now.
+	for _, tc := range testCases {
+		actual, err := NodesConnected(tc.from, tc.to, tc.allowReverse)
+		if err != nil {
+			t.Fatalf("unable to determine node connection: %v", err)
+		}
+		if actual {
+			t.Fatalf("test case %s: nodes connected after commanded to "+
+				"disconnect", tc.name)
+		}
+	}
 }
 
 func testTearDownAll(t *testing.T) {
@@ -492,6 +560,7 @@ var harnessTestCases = []HarnessTestCase{
 	testSendOutputs,
 	testConnectNode,
 	testDisconnectNode,
+	testNodesConnected,
 	testActiveHarnesses,
 	testJoinBlocks,
 	testJoinMempools, // Depends on results of testJoinBlocks

rpctest/utils.go

@@ -166,6 +166,42 @@ func RemoveNode(from *Harness, to *Harness) error {
 	return nil
 }
 
+// NodesConnected verifies whether there is a connection via the p2p interface
+// between the specified nodes. If allowReverse is true, connectivity is also
+// checked in the reverse direction (to->from).
+func NodesConnected(from, to *Harness, allowReverse bool) (bool, error) {
+	peerInfo, err := from.Node.GetPeerInfo()
+	if err != nil {
+		return false, err
+	}
+
+	targetAddr := to.node.config.listen
+	for _, p := range peerInfo {
+		if p.Addr == targetAddr {
+			return true, nil
+		}
+	}
+
+	if !allowReverse {
+		return false, nil
+	}
+
+	// Check in the reverse direction.
+	peerInfo, err = to.Node.GetPeerInfo()
+	if err != nil {
+		return false, err
+	}
+
+	targetAddr = from.node.config.listen
+	for _, p := range peerInfo {
+		if p.Addr == targetAddr {
+			return true, nil
+		}
+	}
+
+	return false, nil
+}
+
 // TearDownAll tears down all active test harnesses.
 func TearDownAll() error {
 	harnessStateMtx.Lock()