libp2p: DisconnectError test (#185)

* DisconnectError tests

pkg/p2p/error.go

@@ -49,9 +49,9 @@ type DisconnectError struct {
 	err error
 }
 
-// Disconnect wraps error and creates a special error that is treated specially
+// NewDisconnectError wraps error and creates a special error that is treated specially
 // by p2p. It causes peer to disconnect.
-func Disconnect(err error) error {
+func NewDisconnectError(err error) error {
 	return &DisconnectError{
 		err: err,
 	}

pkg/p2p/libp2p/libp2p.go

@@ -249,8 +249,7 @@ func (s *Service) AddProtocol(p p2p.ProtocolSpec) (err error) {
 			peerID := streamlibp2p.Conn().RemotePeer()
 			overlay, found := s.peers.overlay(peerID)
 			if !found {
-				// todo: this should never happen, should we disconnect in this case?
-				// todo: test connection close and refactor
+				// todo: this should never happen
 				_ = s.disconnect(peerID)
 				s.logger.Errorf("overlay address for peer %q not found", peerID)
 				return
@@ -280,7 +279,6 @@ func (s *Service) AddProtocol(p p2p.ProtocolSpec) (err error) {
 			if err := ss.Handler(ctx, p2p.Peer{Address: overlay}, stream); err != nil {
 				var e *p2p.DisconnectError
 				if errors.As(err, &e) {
-					// todo: test connection close and refactor
 					_ = s.Disconnect(overlay)
 				}
 

pkg/p2p/libp2p/protocols_test.go

@@ -141,7 +141,36 @@ func TestNewStream_semanticVersioning(t *testing.T) {
 			expectErrNotSupported(t, err)
 		}
 	}
+}
+
+func TestDisconnectError(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	s1, overlay1, cleanup1 := newService(t, libp2p.Options{NetworkID: 1})
+	defer cleanup1()
+
+	s2, overlay2, cleanup2 := newService(t, libp2p.Options{NetworkID: 1})
+	defer cleanup2()
+
+	if err := s1.AddProtocol(newTestProtocol(func(_ context.Context, _ p2p.Peer, _ p2p.Stream) error {
+		return p2p.NewDisconnectError(errors.New("test error"))
+	})); err != nil {
+		t.Fatal(err)
+	}
+
+	addr := serviceUnderlayAddress(t, s1)
+
+	if _, err := s2.Connect(ctx, addr); err != nil {
+		t.Fatal(err)
+	}
+
+	expectPeers(t, s1, overlay2)
 
+	// error is not checked as opening a new stream should cause disconnect from s1 which is async and can make errors in newStream function
+	// it is important to validate that disconnect will happen after NewStream()
+	_, _ = s2.NewStream(ctx, overlay1, nil, testProtocolName, testProtocolVersion, testStreamName)
+	expectPeersEventually(t, s1)
 }
 
 const (