diff options
Diffstat (limited to 'internal/network/network_test.go')
-rw-r--r-- | internal/network/network_test.go | 63 |
1 files changed, 45 insertions, 18 deletions
diff --git a/internal/network/network_test.go b/internal/network/network_test.go index 4f3211c..1bc42fb 100644 --- a/internal/network/network_test.go +++ b/internal/network/network_test.go @@ -5,36 +5,49 @@ import ( "time" "github.com/stretchr/testify/assert" + "go.uber.org/zap" ) // TestPeerToPeerCommunication tests if two peers can communicate. func TestPeerToPeerCommunication(t *testing.T) { // Create a mock of the first peer (peer-1) - peer1IP := "127.0.0.1" - peer1Port := 9001 - peer1 := NewTCPNetwork("peer-1", peer1IP, peer1Port, func() {}) + peer1Opts := TCPNetworkOpts{ + ListenAddr: ":9001", + HandshakeFn: func() error { return nil }, + RetryDelay: time.Second * 2, + Logger: zap.L(), + } + peer1 := NewTCPNetwork("peer-1", peer1Opts) - // Create a mock of the second peer (peer-2) - peer2IP := "127.0.0.1" - peer2Port := 9002 - peer2 := NewTCPNetwork("peer-2", peer2IP, peer2Port, func() {}) - - // Register a message handler on peer-2 to receive the message from peer-1 - peer2.RegisterHandler("chat", func(msg Message) { - assert.Equal(t, "peer-1", msg.Source.ID) - assert.Equal(t, "Hey from peer-1!", string(msg.Payload)) + peer1.RegisterHandler(func(msg Message) { + assert.Equal(t, "Hey from peer-2!", string(msg.Payload)) }) + // Create a mock of the second peer (peer-2) + peer2Opts := TCPNetworkOpts{ + ListenAddr: ":9002", + HandshakeFn: func() error { return nil }, + RetryDelay: time.Second * 2, + Logger: zap.L(), + OnReceiveFn: func(msg Message) { + assert.Equal(t, "Hey from peer-1!", string(msg.Payload)) + }, + } + peer2 := NewTCPNetwork("peer-2", peer2Opts) + // Start the first peer and add the second peer - go peer1.AddPeer("peer-2", peer2IP, peer2Port) - go peer2.AddPeer("peer-1", peer1IP, peer1Port) + go peer1.AddPeer("peer-2", peer2.ListenAddr) + go peer2.AddPeer("peer-1", peer1.ListenAddr) // Wait for the connections to be established // You might need a little more time based on network delay and retry logic time.Sleep(5 * time.Second) // Send a message from peer-1 to peer-2 - err := peer1.Send("peer-2", "chat", []byte("Hey from peer-1!")) + err := peer1.Send("peer-2", []byte("Hey from peer-1!")) + assert.NoError(t, err) + + err = peer2.Send("peer-1", []byte("Hey from peer-2!")) assert.NoError(t, err) // Allow some time for the message to be received and handled @@ -43,10 +56,24 @@ func TestPeerToPeerCommunication(t *testing.T) { // TestSendFailure tests if sending a message fails when no connection exists. func TestSendFailure(t *testing.T) { - peer1 := NewTCPNetwork("peer-1", "127.0.0.1", 9001, func() {}) - _ = NewTCPNetwork("peer-2", "127.0.0.1", 9002, func() {}) + peer1Opts := TCPNetworkOpts{ + ListenAddr: ":9001", + HandshakeFn: DefaultHandshake, + RetryDelay: time.Second * 2, + Logger: zap.L(), + } + peer1 := NewTCPNetwork("peer-1", peer1Opts) + + // Create a mock of the second peer (peer-2) + peer2Opts := TCPNetworkOpts{ + ListenAddr: ":9002", + HandshakeFn: DefaultHandshake, + RetryDelay: time.Second * 2, + Logger: zap.L(), + } + _ = NewTCPNetwork("peer-2", peer2Opts) // Attempt to send a message without establishing a connection first - err := peer1.Send("peer-2", "chat", []byte("Message without connection")) + err := peer1.Send("peer-2", []byte("Message without connection")) assert.Error(t, err, "Expected error when sending to a non-connected peer") } |