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())