Handle MultiselectCommunicatorError

This commit is contained in:
NIC619
2019-09-15 16:58:22 +08:00
parent 879f193aa1
commit eaa74c4e26
2 changed files with 18 additions and 6 deletions

View File

@ -2,7 +2,7 @@ from typing import Dict, Tuple
from libp2p.typing import StreamHandlerFn, TProtocol
from .exceptions import MultiselectError
from .exceptions import MultiselectCommunicatorError, MultiselectError
from .multiselect_communicator_interface import IMultiselectCommunicator
from .multiselect_muxer_interface import IMultiselectMuxer
@ -46,7 +46,10 @@ class Multiselect(IMultiselectMuxer):
# Read and respond to commands until a valid protocol ID is sent
while True:
# Read message
command = await communicator.read()
try:
command = await communicator.read()
except MultiselectCommunicatorError as error:
raise MultiselectError(str(error))
# Command is ls or a protocol
if command == "ls":
@ -76,7 +79,10 @@ class Multiselect(IMultiselectMuxer):
await communicator.write(MULTISELECT_PROTOCOL_ID)
# Read in the protocol ID from other party
handshake_contents = await communicator.read()
try:
handshake_contents = await communicator.read()
except MultiselectCommunicatorError as error:
raise MultiselectError(str(error))
# Confirm that the protocols are the same
if not validate_handshake(handshake_contents):