def test_simple_client_manager_register():
"""Tests if the register method works correctly"""
# Prepare
cid = "1"
bridge = MagicMock()
client = GRPCProxyClient(cid=cid, bridge=bridge)
client_manager = SimpleClientManager()
# Execute
first = client_manager.register(client)
second = client_manager.register(client)
# Assert
assert first
assert not second
assert len(client_manager) == 1
def test_criterion_applied():
"""Test sampling w/ criterion."""
# Prepare
bridge = MagicMock()
client1 = GRPCProxyClient(cid="train_client_1", bridge=bridge)
client2 = GRPCProxyClient(cid="train_client_2", bridge=bridge)
client3 = GRPCProxyClient(cid="test_client_1", bridge=bridge)
client4 = GRPCProxyClient(cid="test_client_2", bridge=bridge)
client_manager = SimpleClientManager()
client_manager.register(client1)
client_manager.register(client2)
client_manager.register(client3)
client_manager.register(client4)
class TestCriterion(Criterion):
"""Criterion to select only test clients."""
def select(self, client: Client) -> bool:
return client.cid.startswith("test_")
# Execute
sampled_clients = client_manager.sample(2, criterion=TestCriterion())
# Assert
assert client3 in sampled_clients
assert client4 in sampled_clients
def test_integration_start_and_shutdown_server():
"""Create server and check if FlowerServiceServicer is returned."""
# Prepare
port = flower_testing.network.unused_tcp_port()
client_manager = SimpleClientManager()
# Execute
server = start_insecure_grpc_server(client_manager=client_manager,
address="[::]",
port=port)
# Teardown
server.stop(1)
def test_integration_connection():
"""Create a server and establish a connection to it.
Purpose of this integration test is to simulate multiple clients
with multiple roundtrips between server and client.
"""
# Prepare
port = flower_testing.network.unused_tcp_port()
server = start_insecure_grpc_server(client_manager=SimpleClientManager(),
address="[::]",
port=port)
# Execute
# Multiple clients in parallel
def run_client() -> int:
messages_received: int = 0
with insecure_grpc_connection(address="[::]", port=port) as conn:
receive, send = conn
# Setup processing loop
while True:
# Block until server responds with a message
server_message = receive()
messages_received += 1
if server_message.HasField("reconnect"):
send(CLIENT_MESSAGE_DISCONNECT)
break
# Process server_message and send client_message...
send(CLIENT_MESSAGE)
return messages_received
results = []
with concurrent.futures.ThreadPoolExecutor() as executor:
futures = [executor.submit(run_client) for _ in range(3)]
concurrent.futures.wait(futures)
for future in futures:
results.append(future.result())
# Assert
for messages_received in results:
assert messages_received == EXPECTED_NUM_SERVER_MESSAGE
# Teardown
server.stop(1)
def test_simple_client_manager_unregister():
"""Tests if the unregister method works correctly"""
# Prepare
cid = "1"
bridge = MagicMock()
client = GrpcClientProxy(cid=cid, bridge=bridge)
client_manager = SimpleClientManager()
client_manager.register(client)
# Execute
client_manager.unregister(client)
# Assert
assert len(client_manager) == 0
def test_criterion_not_applied():
"""Test sampling w/o criterion."""
# Prepare
bridge = MagicMock()
client1 = GRPCProxyClient(cid="train_client_1", bridge=bridge)
client2 = GRPCProxyClient(cid="train_client_2", bridge=bridge)
client3 = GRPCProxyClient(cid="test_client_1", bridge=bridge)
client4 = GRPCProxyClient(cid="test_client_2", bridge=bridge)
client_manager = SimpleClientManager()
client_manager.register(client1)
client_manager.register(client2)
client_manager.register(client3)
client_manager.register(client4)
# Execute
sampled_clients = client_manager.sample(4)
# Assert
assert client1 in sampled_clients
assert client2 in sampled_clients
assert client3 in sampled_clients
assert client4 in sampled_clients