From 33d56c35c7fc8ec0f91c29ebb7603c4fc096c88f Mon Sep 17 00:00:00 2001
From: Gabriel mermelstein <gabrielmermelstein@gmail.com>
Date: Tue, 17 Dec 2024 11:04:01 +0100
Subject: [PATCH] test

---
 waku/nwaku_test.go | 74 ++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 68 insertions(+), 6 deletions(-)

diff --git a/waku/nwaku_test.go b/waku/nwaku_test.go
index 1301fdb..dc23700 100644
--- a/waku/nwaku_test.go
+++ b/waku/nwaku_test.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"slices"
-	"sync"
 	"testing"
 	"time"
 
@@ -461,22 +460,85 @@ func TestRelay(t *testing.T) {
 	defer cancel()
 	senderNode.RelayPublish(ctx, message, pubsubTopic)
 
-	wg := sync.WaitGroup{}
-	wg.Add(1)
-
 	// Wait to receive message
 	select {
 	case envelope := <-receiverNode.node.MsgChan:
 		require.NotNil(t, envelope, "Envelope should be received")
 		require.Equal(t, message.Payload, envelope.Message().Payload, "Received payload should match")
 		require.Equal(t, message.ContentTopic, envelope.Message().ContentTopic, "Content topic should match")
-		wg.Done()
 	case <-time.After(10 * time.Second):
 		t.Fatal("Timeout: No message received within 10 seconds")
 	}
-	wg.Wait()
 
 	// Stop nodes
 	require.NoError(t, senderNode.Stop())
 	require.NoError(t, receiverNode.Stop())
 }
+
+func TestTopicHealth(t *testing.T) {
+	logger, err := zap.NewDevelopment()
+	require.NoError(t, err)
+	clusterId := uint16(16)
+	shardId := uint16(64)
+
+	// start node1
+	wakuConfig1 := WakuConfig{
+		Relay:           true,
+		LogLevel:        "DEBUG",
+		Discv5Discovery: false,
+		ClusterID:       clusterId,
+		Shards:          []uint16{shardId},
+		Discv5UdpPort:   9050,
+		TcpPort:         60050,
+	}
+
+	node1, err := New(&wakuConfig1, logger.Named("node1"))
+	require.NoError(t, err)
+	require.NoError(t, node1.Start())
+	time.Sleep(1 * time.Second)
+
+	// start node2
+	wakuConfig2 := WakuConfig{
+		Relay:           true,
+		LogLevel:        "DEBUG",
+		Discv5Discovery: false,
+		ClusterID:       clusterId,
+		Shards:          []uint16{shardId},
+		Discv5UdpPort:   9051,
+		TcpPort:         60051,
+	}
+	node2, err := New(&wakuConfig2, logger.Named("node2"))
+	require.NoError(t, err)
+	require.NoError(t, node2.Start())
+	time.Sleep(1 * time.Second)
+	multiaddr2, err := node2.ListenAddresses()
+	require.NoError(t, err)
+	require.NotNil(t, multiaddr2)
+
+	// node1 dials node2 so they become peers
+	err = node1.DialPeer(multiaddr2[0])
+	require.NoError(t, err)
+	time.Sleep(1 * time.Second)
+	// Check that both nodes now have one connected peer
+	peerCount1, err := node1.PeerCount()
+	require.NoError(t, err)
+	require.True(t, peerCount1 == 1, "node1 should have 1 peer")
+	peerCount2, err := node2.PeerCount()
+	require.NoError(t, err)
+	require.True(t, peerCount2 == 1, "node2 should have 1 peer")
+
+	// Wait to receive topic health update
+	select {
+	case topicHealth := <-node2.node.TopicHealthChan:
+		require.NotNil(t, topicHealth, "topicHealth should be updated")
+		require.Equal(t, topicHealth.TopicHealth, "MinimallyHealthy", "Topic health should be MinimallyHealthy")
+		require.Equal(t, topicHealth.PubsubTopic, FormatWakuRelayTopic(clusterId, shardId), "PubsubTopic should match configured cluster and shard")
+	case <-time.After(10 * time.Second):
+		t.Fatal("Timeout: No topic health event received within 10 seconds")
+	}
+
+	// Stop nodes
+	require.NoError(t, node1.Stop())
+	require.NoError(t, node2.Stop())
+
+}