summaryrefslogtreecommitdiff
path: root/pkg/p2p/network_test.go
blob: 0b4b6009f8917d40734c6e34c968e98c4e6c9bca (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
package p2p

import (
	"testing"
	"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)
	peer1Opts := TCPNetworkOpts{
		ListenAddr:  ":9001",
		HandshakeFn: DefaultHandshake,
		RetryDelay:  2 * time.Second,
		Logger:      zap.L(),
	}
	peer1 := NewTCPNetwork("peer-1", peer1Opts)
	defer peer1.Close()
	time.Sleep(3 * time.Second)

	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: DefaultHandshake,
		RetryDelay:  2 * time.Second,
		Logger:      zap.L(),
		OnReceiveFn: func(msg Message) {
			assert.Equal(t, "Hey from peer-1!", string(msg.Payload))
		},
	}
	peer2 := NewTCPNetwork("peer-2", peer2Opts)
	defer peer2.Close()
	time.Sleep(3 * time.Second)

	// Start the first peer and add the second peer
	peer1.AddPeer("peer-2", peer2.ListenAddr)
	peer2.AddPeer("peer-1", peer1.ListenAddr)

	// Wait for connections to be established with a timeout
	time.Sleep(5 * time.Second)

	// Send a message from peer-1 to peer-2
	err := peer1.Send("peer-2", []byte("simple-msg"), []byte("Hey from peer-1!"))
	assert.NoError(t, err)

	err = peer2.Send("peer-1", []byte("simple-msg"), []byte("Hey from peer-2!"))
	assert.NoError(t, err)

	// Allow some time for the message to be received and handled
	time.Sleep(2 * time.Second)
}

// TestSendFailure tests if sending a message fails when no connection exists.
func TestSendFailure(t *testing.T) {
	peer1Opts := TCPNetworkOpts{
		ListenAddr:  ":9003",
		HandshakeFn: DefaultHandshake,
		RetryDelay:  time.Second * 2,
		Logger:      zap.L(),
	}
	peer1 := NewTCPNetwork("peer-1", peer1Opts)
	defer peer1.Close()

	// Create a mock of the second peer (peer-2) - but don't add it to peer1
	peer2Opts := TCPNetworkOpts{
		ListenAddr:  ":9004",
		HandshakeFn: DefaultHandshake,
		RetryDelay:  time.Second * 2,
		Logger:      zap.L(),
	}
	peer2 := NewTCPNetwork("peer-2", peer2Opts)
	defer peer2.Close()

	// Attempt to send a message without establishing a connection first
	err := peer1.Send("peer-2", []byte("msg"), []byte("Message without connection"))
	assert.Error(t, err, "Expected error when sending to a non-connected peer")
}