def client_disconnect_to_proto(reason: str) -> ClientMessage.Disconnect:
reason_proto = Reason.UNKNOWN
if reason == "RECONNECT":
reason_proto = Reason.RECONNECT
elif reason == "POWER_DISCONNECTED":
reason_proto = Reason.POWER_DISCONNECTED
elif reason == "WIFI_UNAVAILABLE":
reason_proto = Reason.WIFI_UNAVAILABLE
return ClientMessage.Disconnect(reason=reason_proto)
import grpc
import flower_testing
from flower.client_manager import SimpleClientManager
from flower.grpc_client.connection import insecure_grpc_connection
from flower.grpc_server.grpc_server import start_insecure_grpc_server
from flower.proto.transport_pb2 import ClientMessage, ServerMessage
EXPECTED_NUM_SERVER_MESSAGE = 10
SERVER_MESSAGE = ServerMessage()
SERVER_MESSAGE_RECONNECT = ServerMessage(reconnect=ServerMessage.Reconnect())
CLIENT_MESSAGE = ClientMessage()
CLIENT_MESSAGE_DISCONNECT = ClientMessage(
disconnect=ClientMessage.Disconnect())
def mock_join( # type: ignore # pylint: disable=invalid-name
_self,
request_iterator: Iterator[ClientMessage],
_context: grpc.ServicerContext,
) -> Iterator[ServerMessage]:
"""Serve as mock for the Join method of class FlowerServiceServicer."""
counter = 0
while True:
counter += 1
if counter < EXPECTED_NUM_SERVER_MESSAGE:
yield SERVER_MESSAGE