From 0a61e83a154efb9d0955d1db3558111dcdb21771 Mon Sep 17 00:00:00 2001 From: Daniel Goldberg Date: Fri, 27 Sep 2019 16:57:36 +0300 Subject: [PATCH] Add chunking test and some basic docs --- .../model/victim_host_generator_test.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/model/victim_host_generator_test.py b/monkey/infection_monkey/model/victim_host_generator_test.py index 849196ff6..2414bd794 100644 --- a/monkey/infection_monkey/model/victim_host_generator_test.py +++ b/monkey/infection_monkey/model/victim_host_generator_test.py @@ -3,16 +3,25 @@ from infection_monkey.model.victim_host_generator import VictimHostGenerator from common.network.network_range import CidrRange, SingleIpRange -class TestPayload(TestCase): +class VictimHostGeneratorTester(TestCase): def setUp(self): - self.test_ranges = [CidrRange("10.0.0.0/28", False), + self.test_ranges = [CidrRange("10.0.0.0/28", False), # this gives us 15 hosts SingleIpRange('41.50.13.37'), SingleIpRange('localhost') ] self.generator = VictimHostGenerator(self.test_ranges, '10.0.0.1') self.generator._ip_addresses = [] # test later on + def test_chunking(self): + chunk_size = 3 + # current test setup is 15+1+1-1 hosts + victims = self.generator.generate_victims(chunk_size) + for i in range(5): # quickly check the equally sided chunks + self.assertEqual(len(victims.next()), chunk_size) + victim_chunk_last = victims.next() + self.assertEqual(len(victim_chunk_last), 1) + def test_remove_blocked_ip(self): victims = list(self.generator.generate_victims_from_range(self.test_ranges[0])) self.assertEqual(len(victims), 14) # 15 minus the 1 we blocked