From a27a817d50d76bffeeed1f4d83a10168944bfebe Mon Sep 17 00:00:00 2001
From: mhchia <kevin.mh.chia@gmail.com>
Date: Fri, 20 Sep 2019 16:17:13 +0800
Subject: [PATCH] Fix tests

---
 libp2p/host/basic_host.py                   | 6 +++---
 libp2p/network/swarm.py                     | 6 +++++-
 tests/examples/test_chat.py                 | 4 ++--
 tests/protocol_muxer/test_protocol_muxer.py | 6 +++---
 4 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/libp2p/host/basic_host.py b/libp2p/host/basic_host.py
index d434c419..126359aa 100644
--- a/libp2p/host/basic_host.py
+++ b/libp2p/host/basic_host.py
@@ -1,21 +1,21 @@
 import asyncio
-from typing import List, Sequence
 import logging
+from typing import List, Sequence
 
 import multiaddr
 
+from libp2p.host.exceptions import StreamFailure
 from libp2p.network.network_interface import INetwork
 from libp2p.network.stream.net_stream_interface import INetStream
 from libp2p.peer.id import ID
 from libp2p.peer.peerinfo import PeerInfo
 from libp2p.peer.peerstore_interface import IPeerStore
+from libp2p.protocol_muxer.exceptions import MultiselectClientError, MultiselectError
 from libp2p.protocol_muxer.multiselect import Multiselect
 from libp2p.protocol_muxer.multiselect_client import MultiselectClient
 from libp2p.protocol_muxer.multiselect_communicator import MultiselectCommunicator
 from libp2p.routing.kademlia.kademlia_peer_router import KadmeliaPeerRouter
 from libp2p.typing import StreamHandlerFn, TProtocol
-from libp2p.protocol_muxer.exceptions import MultiselectClientError, MultiselectError
-from libp2p.host.exceptions import StreamFailure
 
 from .host_interface import IHost
 
diff --git a/libp2p/network/swarm.py b/libp2p/network/swarm.py
index 47d76ebc..7a1b0990 100644
--- a/libp2p/network/swarm.py
+++ b/libp2p/network/swarm.py
@@ -10,7 +10,11 @@ from libp2p.peer.peerstore import PeerStoreError
 from libp2p.peer.peerstore_interface import IPeerStore
 from libp2p.routing.interfaces import IPeerRouting
 from libp2p.stream_muxer.abc import IMuxedConn
-from libp2p.transport.exceptions import MuxerUpgradeFailure, SecurityUpgradeFailure
+from libp2p.transport.exceptions import (
+    MuxerUpgradeFailure,
+    OpenConnectionError,
+    SecurityUpgradeFailure,
+)
 from libp2p.transport.listener_interface import IListener
 from libp2p.transport.transport_interface import ITransport
 from libp2p.transport.upgrader import TransportUpgrader
diff --git a/tests/examples/test_chat.py b/tests/examples/test_chat.py
index e2aa71de..99e4380d 100644
--- a/tests/examples/test_chat.py
+++ b/tests/examples/test_chat.py
@@ -2,7 +2,7 @@ import asyncio
 
 import pytest
 
-from libp2p.network.exceptions import SwarmException
+from libp2p.host.exceptions import StreamFailure
 from libp2p.peer.peerinfo import info_from_p2p_addr
 from tests.utils import set_up_nodes_by_transport_opt
 
@@ -84,7 +84,7 @@ async def no_common_protocol(host_a, host_b):
     host_a.set_stream_handler(PROTOCOL_ID, stream_handler)
 
     # try to creates a new new with a procotol not known by the other host
-    with pytest.raises(SwarmException):
+    with pytest.raises(StreamFailure):
         await host_b.new_stream(host_a.get_id(), ["/fakeproto/0.0.1"])
 
 
diff --git a/tests/protocol_muxer/test_protocol_muxer.py b/tests/protocol_muxer/test_protocol_muxer.py
index 4e58e5b6..6ff401cf 100644
--- a/tests/protocol_muxer/test_protocol_muxer.py
+++ b/tests/protocol_muxer/test_protocol_muxer.py
@@ -1,6 +1,6 @@
 import pytest
 
-from libp2p.network.exceptions import SwarmException
+from libp2p.host.exceptions import StreamFailure
 from tests.utils import echo_stream_handler, set_up_nodes_by_transport_opt
 
 # TODO: Add tests for multiple streams being opened on different
@@ -47,7 +47,7 @@ async def test_single_protocol_succeeds():
 
 @pytest.mark.asyncio
 async def test_single_protocol_fails():
-    with pytest.raises(SwarmException):
+    with pytest.raises(StreamFailure):
         await perform_simple_test("", ["/echo/1.0.0"], ["/potato/1.0.0"])
 
     # Cleanup not reached on error
@@ -77,7 +77,7 @@ async def test_multiple_protocol_second_is_valid_succeeds():
 async def test_multiple_protocol_fails():
     protocols_for_client = ["/rock/1.0.0", "/foo/1.0.0", "/bar/1.0.0"]
     protocols_for_listener = ["/aspyn/1.0.0", "/rob/1.0.0", "/zx/1.0.0", "/alex/1.0.0"]
-    with pytest.raises(SwarmException):
+    with pytest.raises(StreamFailure):
         await perform_simple_test("", protocols_for_client, protocols_for_listener)
 
     # Cleanup not reached on error