async def test_key_manager():
key_manager = GroupKeyManager(hivemind.DHT(start=True), endpoint='localhvost',
prefix='test_averaging', initial_group_bits='10110',
target_group_size=2)
t = hivemind.get_dht_time()
key = key_manager.current_key
await key_manager.declare_averager(key, 'localhvost', expiration_time=t + 60)
await key_manager.declare_averager(key, 'localhvost2', expiration_time=t + 61)
q1 = await key_manager.get_averagers(key, only_active=True)
await key_manager.declare_averager(key, 'localhvost', expiration_time=t + 66)
q2 = await key_manager.get_averagers(key, only_active=True)
await key_manager.declare_averager(key, 'localhvost2', expiration_time=t + 61, looking_for_group=False)
q3 = await key_manager.get_averagers(key, only_active=True)
q4 = await key_manager.get_averagers(key, only_active=False)
q5 = await key_manager.get_averagers('nonexistent_key.0b0101', only_active=False)
assert len(q1) == 2 and ('localhvost', t + 60) in q1 and ('localhvost2', t + 61) in q1
assert len(q2) == 2 and ('localhvost', t + 66) in q2 and ('localhvost2', t + 61) in q2
assert len(q3) == 1 and ('localhvost', t + 66) in q3
assert len(q4) == 2 and ('localhvost', t + 66) in q4 and ('localhvost2', t + 61) in q2
assert len(q5) == 0
def __init__(self, endpoint: Endpoint, schema_hash: bytes, dht: DHT, *,
prefix: str, target_group_size: int, min_group_size: int,
request_timeout: float, client_mode: bool, initial_group_bits: Optional[str] = None,
averaging_expiration: float = 15):
assert '.' not in prefix, "group prefix must be a string without ."
if request_timeout is None or request_timeout >= averaging_expiration:
logger.warning("It is recommended to use request_timeout smaller than averaging_expiration. Otherwise,"
"matchmaking can cause deadlocks in some rare cases. Please see Matchmaking docstring.")
super().__init__()
self.endpoint, self.schema_hash = endpoint, schema_hash
self.group_key_manager = GroupKeyManager(dht, endpoint, prefix, initial_group_bits, target_group_size)
self.target_group_size, self.min_group_size = target_group_size, min_group_size
self.averaging_expiration, self.request_timeout = averaging_expiration, request_timeout
self.client_mode = client_mode
self.lock_looking_for_group = asyncio.Lock()
self.lock_request_join_group = asyncio.Lock()
self.follower_was_discarded = asyncio.Event()
self.was_accepted_to_group = asyncio.Event()
self.assembled_group = asyncio.Future()
self.current_leader: Optional[Endpoint] = None # iff i am a follower, this is a link to my current leader
self.current_followers: Dict[Endpoint, averaging_pb2.JoinRequest] = {} # my current followers excluding myself
self.potential_leaders = PotentialLeaders(endpoint, averaging_expiration, target_group_size)
self.data_for_gather: Optional[bytes] = None