async def test_candidate_sends_votes(self, unused_tcp_port_factory): port_c, port_f = unused_tcp_port_factory(), unused_tcp_port_factory() self.communicator_c = HTTPCommunicator('candidate', port_c) self.communicator_f = HTTPCommunicator('follower', port_f) await self.communicator_f.start() self.raft_c = RaftConsensus( self.communicator_c, Registry([{ 'host': 'localhost', 'port': port_f, 'identifier': 'follower' }])) self.raft_f = RaftConsensus( self.communicator_f, Registry([{ 'host': 'localhost', 'port': port_c, 'identifier': 'candidate' }])) self.raft_c.state = PeerState.CANDIDATE self.raft_c.term = 4 on_rv = Mock() on_rv.return_value = True, {"vote_granted": True} self.communicator_f.set_on_request_vote(on_rv) await self.raft_c.process_state() on_rv.assert_called_with({'identifier': 'candidate', 'term': 4})
async def test_leader_gives_up_on_followers_long_heartbeat( self, unused_tcp_port_factory): port_l, port_f = unused_tcp_port_factory(), unused_tcp_port_factory() self.communicator_l = HTTPCommunicator('leader', port_l) self.communicator_f = HTTPCommunicator('follower', port_f) await self.communicator_f.start() self.raft_l = RaftConsensus( self.communicator_l, Registry([{ 'host': 'localhost', 'port': port_f, 'identifier': 'follower' }])) self.raft_f = RaftConsensus( self.communicator_f, Registry([{ 'host': 'localhost', 'port': port_l, 'identifier': 'leader' }])) self.raft_l.state = PeerState.LEADER async def long_call(data): await asyncio.sleep(0.5) self.communicator_f.set_on_heartbeat(long_call) t_start = time.time() await self.raft_l.process_state() duration = time.time() - t_start - 0.050 assert duration <= 0.15
async def test_old_candidate_rejected_by_newer_term( self, unused_tcp_port_factory): port_c, port_f = unused_tcp_port_factory(), unused_tcp_port_factory() self.communicator_c = HTTPCommunicator('candidate', port_c) self.communicator_f = HTTPCommunicator('leader', port_f) await self.communicator_f.start() self.raft_c = RaftConsensus( self.communicator_c, Registry([{ 'host': 'localhost', 'port': port_f, 'identifier': 'follower' }])) self.raft_f = RaftConsensus( self.communicator_f, Registry([{ 'host': 'localhost', 'port': port_c, 'identifier': 'candidate' }])) self.raft_c.state = PeerState.CANDIDATE self.raft_c.term = 3 self.raft_f.term = 4 await self.raft_c.process_state() assert self.raft_c.state == PeerState.CANDIDATE
async def test_invalid_heartbeat_does_not_cancel_election( self, unused_tcp_port_factory): port_c, port_l = unused_tcp_port_factory(), unused_tcp_port_factory() self.communicator_c = HTTPCommunicator('candidate', port_c) self.communicator_l = HTTPCommunicator('leader', port_l) await self.communicator_l.start() self.raft_c = RaftConsensus( self.communicator_c, Registry([{ 'host': 'localhost', 'port': port_l, 'identifier': 'leader' }])) self.raft_l = RaftConsensus( self.communicator_l, Registry([{ 'host': 'localhost', 'port': port_c, 'identifier': 'candidate' }])) self.raft_c.state = PeerState.CANDIDATE self.raft_c.term = 4 self.raft_l.term = 4 async def rv(data): self.raft_c.on_heartbeat({ 'identifier': 'leader', 'term': self.raft_c.term - 1 }) self.communicator_l.set_on_request_vote(rv) await self.raft_c.process_state() assert self.raft_c.state != PeerState.FOLLOWER
async def test_candidate_only_votes_once(self, unused_tcp_port_factory): port_1, port_2 = unused_tcp_port_factory(), unused_tcp_port_factory() self.communicator_1 = HTTPCommunicator('candidate_1', port_1) self.communicator_2 = HTTPCommunicator('candidate_2', port_2) await self.communicator_1.start() await self.communicator_2.start() self.raft_1 = RaftConsensus( self.communicator_1, Registry([{ 'host': 'localhost', 'port': port_1, 'identifier': 'candidate_1' }])) self.raft_2 = RaftConsensus( self.communicator_2, Registry([{ 'host': 'localhost', 'port': port_2, 'identifier': 'candidate_2' }])) self.raft_1.state = PeerState.CANDIDATE self.raft_2.state = PeerState.CANDIDATE async def rv(): return True, {"vote_granted": True} self.communicator_2.set_on_request_vote(rv) await self.raft_1.process_state() await self.raft_2.process_state() assert self.raft_1.state == PeerState.CANDIDATE assert self.raft_2.state == PeerState.CANDIDATE
async def test_leader_sets_followers_heartbeat_flag( self, unused_tcp_port_factory): port_l, port_f = unused_tcp_port_factory(), unused_tcp_port_factory() self.communicator_l = HTTPCommunicator('leader', port_l) self.communicator_f = HTTPCommunicator('follower', port_f) await self.communicator_f.start() self.raft_l = RaftConsensus( self.communicator_l, Registry([{ 'host': 'localhost', 'port': port_f, 'identifier': 'follower' }])) self.raft_f = RaftConsensus( self.communicator_f, Registry([{ 'host': 'localhost', 'port': port_l, 'identifier': 'leader' }])) self.raft_l.state = PeerState.LEADER assert self.raft_f.got_heartbeat.is_set() is False await self.raft_l.process_state() assert self.raft_f.got_heartbeat.is_set() await self.raft_f.process_state() assert self.raft_f.got_heartbeat.is_set() is False
async def test_creates_requestor(self): """Test that a requester is initialized""" # Setup self.communicator = HTTPCommunicator('a', 7000) # Assert assert self.communicator._requester is not None
async def test_on_heartbeat_rejects_invalid_data(self, unused_tcp_port): """Test that the data parameter rejects non None or dict data.""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() # Assert with pytest.raises(ValueError): await self.communicator.send_heartbeat('localhost', unused_tcp_port, 5)
async def test_ping_returns_false(self, unused_tcp_port): """Test that a ping doesn't come back""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() # Act result = await self.communicator.ping('localhost', 0) # Assert assert result is False
async def test_follower_will_vote_for_newer_term(self, unused_tcp_port_factory): port_1, port_2, port_3 = unused_tcp_port_factory( ), unused_tcp_port_factory(), unused_tcp_port_factory() self.communicator_1 = HTTPCommunicator('candidate_1', port_1) self.communicator_2 = HTTPCommunicator('candidate_2', port_2) self.communicator_3 = HTTPCommunicator('candidate_3', port_3) await self.communicator_1.start() await self.communicator_2.start() await self.communicator_3.start() self.raft_1 = RaftConsensus( self.communicator_1, Registry([{ 'host': 'localhost', 'port': port_1, 'identifier': 'candidate_1' }])) self.raft_2 = RaftConsensus( self.communicator_2, Registry([{ 'host': 'localhost', 'port': port_2, 'identifier': 'candidate_2' }])) self.raft_3 = RaftConsensus( self.communicator_3, Registry([{ 'host': 'localhost', 'port': port_3, 'identifier': 'candidate_3' }])) self.raft_1.term = 87 self.raft_1.state = PeerState.CANDIDATE self.raft_1.has_voted_in_term = True self.raft_2.term = 87 self.raft_2.state = PeerState.CANDIDATE self.raft_2.has_voted_in_term = True self.raft_3.term = 86 self.raft_3.state = PeerState.FOLLOWER # Simulate raft_1 sending a request vote data_1 = {"identifier": "candidate_1", "term": self.raft_1.term} assert self.raft_2.on_request_vote(data_1) == (True, { "vote_granted": False }) assert self.raft_3.on_request_vote(data_1) == (True, { "vote_granted": True }) assert self.raft_3.term == 87 assert self.raft_3.state == PeerState.FOLLOWER assert self.raft_3.has_voted_in_term
async def test_register_returns_true(self, unused_tcp_port): """Test that registering returns true.""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() # Act status = await self.communicator.register_with('localhost', unused_tcp_port) # Assert assert status is True
async def test_heartbeat_returns_true(self): """Test that sending a heartbeat returns with response code 200""" # Setup self.communicator = HTTPCommunicator('a', 7000) await self.communicator.start() # Act status, data = await self.communicator.send_heartbeat( 'localhost', 7000, {'identifier': 'a'}) # Assert assert status is True assert data is None
async def test_ping_timeout(self, unused_tcp_port): """Test that a ping will timeout""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() # Act status = await self.communicator.ping('localhost', unused_tcp_port, timeout=0) # Assert assert status is False
async def test_on_register_executes(self, unused_tcp_port): """Test that our register callback is called with the right data""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() on_register = Mock() self.communicator.set_on_register(on_register) # Act await self.communicator.register_with('localhost', unused_tcp_port) # Assert on_register.assert_called_with('localhost', unused_tcp_port, 'a')
async def test_heartbeat_returns_false_on_error(self, unused_tcp_port): """Test that nothing returns on a connection error""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() # Act status, data = await self.communicator.send_heartbeat( 'localhost', 0, {}) # Assert assert status is False assert data is None
async def test_on_heartbeat_transforms_none_data(self, unused_tcp_port): """Test that the data parameter is an empty dict when None is supplied""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() cb = Mock() self.communicator.set_on_heartbeat(cb) # Act await self.communicator.send_heartbeat('localhost', unused_tcp_port, None) # Assert cb.assert_called_with({})
async def test_request_vote_passes_data_to_callback(self, unused_tcp_port): """Test that request vote passed the data to the defined callback""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() on_request_vote = Mock() self.communicator.set_on_request_vote(on_request_vote) data = {'test': True} # Act await self.communicator.request_vote('localhost', unused_tcp_port, data) # Assert on_request_vote.assert_called_with(data)
async def test_on_register_can_return_error(self, unused_tcp_port): """Test that our register callback can force an error""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() on_register = Mock() on_register.return_value = (False, "Error!") self.communicator.set_on_register(on_register) # Act status = await self.communicator.register_with('localhost', unused_tcp_port) # Assert assert status is False
async def test_on_heartbeat_can_return_error(self, unused_tcp_port): """Test that our function can cause an error to be returned""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() on_heartbeat = Mock() on_heartbeat.return_value = (False, {"status", "error!"}) self.communicator.set_on_heartbeat(on_heartbeat) # Act status, data = await self.communicator.send_heartbeat( 'localhost', unused_tcp_port, {}) # Assert assert status is False assert data is None
async def test_request_vote_returns_400_with_no_callback( self, unused_tcp_port): """Test that request vote returns a 400 status when no callback is defined""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() self.communicator.set_on_request_vote(None) # Act res = await self.communicator._requester.post('localhost', unused_tcp_port, '/raft/request_vote', {'test': 123}) # Assert assert res.status == 400
async def test_request_vote_returns_on_error(self, unused_tcp_port): """Test that request vote will return nothing on error""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() on_request_vote = Mock() self.communicator.set_on_request_vote(on_request_vote) data = {'test': True} # Act status, data = await self.communicator.request_vote( 'localhost', 0, data) # Assert assert status is False assert data is None
async def test_on_heartbeat_executes(self, unused_tcp_port): """Test that our function is called each time we receive a heartbeat""" # Setup self.communicator = HTTPCommunicator('a', unused_tcp_port) await self.communicator.start() on_heartbeat = Mock() on_heartbeat.return_value = True, {} self.communicator.set_on_heartbeat(on_heartbeat) data = {'identifier': 'a', 'term': 1} # Act await self.communicator.send_heartbeat('localhost', unused_tcp_port, data) # Assert on_heartbeat.assert_called_with(data)
class QCluster(object): def __init__(self, identifier='', listen_host='localhost', listen_port=0, peers=[]): self.identifier = identifier self.listen_host = listen_host self.listen_port = int(listen_port) # MARK: Setup the communication module event_loop = asyncio.get_event_loop() self.communicator = HTTPCommunicator(self.identifier, listen_host=self.listen_host, listen_port=self.listen_port) self.registry = Registry(peers) self.raft = RaftConsensus(self.communicator, self.registry) event_loop.create_task(self.communicator.start()) event_loop.create_task(self.raft.start()) def is_leader(self): return self.raft.state == PeerState.LEADER def get_leader_info(self): if self.raft.known_leader is not None: return self.registry.get_peer_by_identifier(self.raft.known_leader) else: return None
async def test_heartbeat_follower_stays_follower(self, unused_tcp_port): self.communicator = HTTPCommunicator('a', unused_tcp_port) self.registry = Registry([]) self.raft = RaftConsensus(self.communicator, self.registry) self.raft.got_heartbeat.set() await self.raft.process_state() assert self.raft.state == PeerState.FOLLOWER
async def test_is_follower_returns_false_when_not_folower( self, unused_tcp_port): self.communicator = HTTPCommunicator('a', unused_tcp_port) self.raft = RaftConsensus(self.communicator, Registry([])) self.raft.state = PeerState.LEADER assert self.raft.is_follower() is False self.raft.state = PeerState.CANDIDATE assert self.raft.is_follower() is False
async def test_get_timeout_returns_valid_default_timeout_value( self, unused_tcp_port): self.communicator = HTTPCommunicator('a', unused_tcp_port) self.registry = Registry([]) self.raft = RaftConsensus(self.communicator, self.registry) t = self.raft.get_timeout() assert 0.150 <= t <= 0.300
async def test_is_candidate_returns_false_when_not_candidate( self, unused_tcp_port): self.communicator = HTTPCommunicator('a', unused_tcp_port) self.raft = RaftConsensus(self.communicator, Registry([])) self.raft.state = PeerState.FOLLOWER assert self.raft.is_candidate() is False self.raft.state = PeerState.LEADER assert self.raft.is_candidate() is False
async def test_no_heartbeat_follower_becomes_candidate( self, unused_tcp_port): self.communicator = HTTPCommunicator('a', unused_tcp_port) self.registry = Registry([]) self.raft = RaftConsensus(self.communicator, self.registry) await self.raft.process_state() assert self.raft.term == 1 assert self.raft.state == PeerState.CANDIDATE
async def test_valid_heartbeat_sets_got_heartbeat(self, unused_tcp_port): self.communicator = HTTPCommunicator('a', unused_tcp_port) self.registry = Registry([]) self.raft = RaftConsensus(self.communicator, self.registry) data = {'identifier': '1', 'term': self.raft.term} self.raft.on_heartbeat(data) assert self.raft.got_heartbeat.is_set()
def __init__(self, identifier='', listen_host='localhost', listen_port=0, peers=[]): self.identifier = identifier self.listen_host = listen_host self.listen_port = int(listen_port) # MARK: Setup the communication module event_loop = asyncio.get_event_loop() self.communicator = HTTPCommunicator(self.identifier, listen_host=self.listen_host, listen_port=self.listen_port) self.registry = Registry(peers) self.raft = RaftConsensus(self.communicator, self.registry) event_loop.create_task(self.communicator.start()) event_loop.create_task(self.raft.start())