def test_discovery_inc_stage(self, mock_file):
"""
The discovery stage needs to finish to produce the list of live
targets for the other stages.
"""
mock_file.side_effect = self.side_effect
context = Context(self.mock_server_config)
self.assertEqual(3, len(context.stage_list))
context.pop("127.0.0.1:1010")
context.pop("127.0.0.2:1010")
self.assertIsNone(context.pop("127.0.0.3:1010"))
def test_save_context(self):
with patch('builtins.open', mock_open()) as mopen:
handle = mopen.return_value
ctx = Context(self.config)
self.config.save_context(ctx)
mopen.assert_any_call('data/run/current.trace', "wb")
self.assertEqual(1, handle.write.call_count)
handle.write.assert_any_call(pickle.dumps(ctx))
def test_context_resume(self, mock_file):
mock_file.side_effect = self.side_effect
agent1 = "127.0.0.1:1010"
agent2 = "127.0.0.2:1010"
context = Context(self.mock_server_config)
task1 = context.pop(agent1)
task2 = context.pop(agent2)
context.running(agent1)
self.mock_server_config.save_context(self.mock_server_config, context)
mock_file.assert_any_call('fake/run/current.trace', "wb")
restored_ctx = Context.create(self.mock_server_config)
self.assertEqual(context.cstage_name, restored_ctx.cstage_name)
self.assertEqual(len(context.stage_list), len(restored_ctx.stage_list))
self.assertEqual(0, len(restored_ctx.active))
self.assertEqual(context.reports_path, restored_ctx.reports_path)
self.assertEqual(2, len(restored_ctx.pending))
rtask1 = restored_ctx.pop(agent1)
self.assertEqual(task1, rtask1)
self.assertEqual(1, len(restored_ctx.active))
rtask2 = restored_ctx.pop(agent2)
self.assertEqual(task2, rtask2)
self.assertEqual(2, len(restored_ctx.active))
self.assertIsNotNone(restored_ctx._lock)
restored_ctx.completed(agent1)
restored_ctx.completed(agent2)
_ = restored_ctx.pop(agent1)
restored_ctx.completed(agent1)
restored_ctx.completed(agent2)
nstages, pending, completion = restored_ctx.ctx_status()[0]
self.assertEqual(3, nstages)
self.assertEqual(0, pending)
self.assertEqual("66.67%", completion)
self.assertEqual("stage1", restored_ctx.cstage_name)
def test_status_update(self, mock_file):
mock_file.side_effect = self.side_effect
agent = "127.0.0.1:1010"
context = Context(self.mock_server_config)
_ = context.pop(agent)
# the completed status is short lived as it's deleted after the
# task is updated.
test_cases = [
("running", STATUS.RUNNING),
("downloading", STATUS.DOWNLOADING),
("interrupted", STATUS.INTERRUPTED),
]
for name, expected in test_cases:
method = getattr(context, name)
method(agent)
task = context.active.get(agent)
if name == "interrupted":
self.assertEqual(None, task)
self.assertEqual(1, len(context.pending))
task = context.pending.pop(0)
self.assertEqual(expected, task.status)
def test_interruption(self, mock_file):
mock_file.side_effect = self.side_effect
agent = "127.0.0.1:1010"
context = Context(self.mock_server_config)
_ = context.pop(agent)
context.interrupted(agent)
self.assertEqual(1, len(context.pending))
_ = context.pop(agent)
self.assertEqual(0, len(context.pending))
def test_create(self, os_stat):
os_stat.return_value = Mock(st_size=0)
ctx = Context.create(self.mock_server_config)
self.assertIsNotNone(ctx)
def test_context_full_flow(self, mock_file):
mock_file.side_effect = self.side_effect
context = Context(self.mock_server_config)
self.assertEqual(3, len(context.stage_list))
expected_active_status = [('discovery', 2, 2, "100.00%"),
('stage1', 1, 1, "100.00%"),
('stage2', 1, 0, "0.00%")]
expected_task_status = [
('127.0.0.1:1010', 'stage1', 'SCHEDULED', '172.16.71.132/31'),
('127.0.0.2:1010', 'stage2', 'SCHEDULED', '172.16.71.132/31')
]
agent = "127.0.0.1:1010"
# starts with stage 1 discovery
task_data = context.pop(agent)
task = context.active.get(agent)
nactive = len(context.active)
self.check_tasks(task_data, task, "discovery", nactive,
"172.16.71.132")
context.completed(agent)
# stage 1 discovery target 2
task_data = context.pop(agent)
task = context.active.get(agent)
nactive = len(context.active)
self.check_tasks(task_data, task, "discovery", nactive,
"172.16.71.133")
context.completed(agent)
# stage 1
task_data = context.pop(agent)
task = context.active.get(agent)
nactive = len(context.active)
self.check_tasks(task_data, task, "stage1", nactive,
"172.16.71.132/31")
# if another agent pulls now should be 2 active and stage 2 should
# start
task_data = context.pop("127.0.0.2:1010")
task = context.active.get("127.0.0.2:1010")
nactive = len(context.active)
self.check_tasks(task_data, task, "stage2", nactive,
"172.16.71.132/31", 2)
task_status = context.tasks_status()
self.assertEqual(expected_task_status, task_status)
context.completed(agent)
nstages, pending, completion = context.ctx_status()[0]
self.assertEqual(3, nstages)
self.assertEqual(0, pending)
self.assertEqual("66.67%", completion)
active_status = context.active_stages_status()
self.assertEqual(expected_active_status, active_status)
context.completed("127.0.0.2:1010")
nstages, pending, completion = context.ctx_status()[0]
self.assertEqual(3, nstages)
self.assertEqual(0, pending)
self.assertEqual("100.00%", completion)
self.assertIsNone(context.pop(agent))
def __init__(self, *args, options, **kwargs):
self._terminate = threading.Event()
self.options = options
self.ctx = Context.create(options)
super().__init__(*args, **kwargs)