def test_increase_number_of_workers(self):
worker_pool = Pool(processes=2,
worker_names='WorkerThread',
maxtasksperchild=3)
self.assertEqual(worker_pool.get_worker_count(), 2)
def noop():
return 1 + 2
for _ in xrange(12):
result = worker_pool.apply_async(func=noop)
self.assertEqual(result.get(), 3)
self.assertEqual(worker_pool.get_worker_count(), 2)
worker_pool.set_worker_count(4)
# It takes some time...
self.assertEqual(worker_pool.get_worker_count(), 2)
for _ in xrange(12):
result = worker_pool.apply_async(func=noop)
self.assertEqual(result.get(), 3)
self.assertEqual(worker_pool.get_worker_count(), 4)
worker_pool.terminate()
worker_pool.join()
def test_multiple_append_uniq_group(self):
def multi_append():
for i in xrange(InfoSet.MAX_INFO_INSTANCES * 2):
vuln = MockVuln()
kb.append_uniq_group('a', 'b', vuln, group_klass=MockInfoSetTrue)
info_set_list = kb.get('a', 'b')
self.assertEqual(len(info_set_list), 1)
info_set = info_set_list[0]
self.assertEqual(len(info_set.infos), InfoSet.MAX_INFO_INSTANCES)
return True
pool = Pool(2)
r1 = pool.apply_async(multi_append)
r2 = pool.apply_async(multi_append)
r3 = pool.apply_async(multi_append)
self.assertTrue(r1.get())
self.assertTrue(r2.get())
self.assertTrue(r3.get())
pool.terminate()
pool.join()
def test_multiple_append_uniq_group(self):
def multi_append():
for i in xrange(InfoSet.MAX_INFO_INSTANCES * 2):
vuln = MockVuln()
kb.append_uniq_group('a',
'b',
vuln,
group_klass=MockInfoSetTrue)
info_set_list = kb.get('a', 'b')
self.assertEqual(len(info_set_list), 1)
info_set = info_set_list[0]
self.assertEqual(len(info_set.infos), InfoSet.MAX_INFO_INSTANCES)
return True
pool = Pool(2)
r1 = pool.apply_async(multi_append)
r2 = pool.apply_async(multi_append)
r3 = pool.apply_async(multi_append)
self.assertTrue(r1.get())
self.assertTrue(r2.get())
self.assertTrue(r3.get())
pool.terminate()
pool.join()
def test_pickleable_shells(self):
pool = Pool(1)
xurllib = ExtendedUrllib()
original_shell = Shell(MockVuln(), xurllib, pool)
kb.append('a', 'b', original_shell)
unpickled_shell = kb.get('a', 'b')[0]
self.assertEqual(original_shell, unpickled_shell)
self.assertEqual(unpickled_shell.worker_pool, None)
self.assertEqual(unpickled_shell._uri_opener, None)
pool.terminate()
pool.join()
xurllib.end()
def get_fingerprint(url, threads):
pool = Pool(worker_names='HMap',
maxtasksperchild=2,
processes=threads,
max_queued_tasks=5)
tests = {
basic_get, basic_options, unknown_method, unauthorized_activity,
nonexistant_object, malformed_method_line, long_url_ranges,
long_default_ranges, many_header_ranges, large_header_ranges,
unavailable_accept, fake_content_length
}
for test in tests:
pool.apply_async(func=test, args=(url, ))
pool.close()
pool.join()
pool.terminate()
fingerprint['SYNTACTIC']['HEADER_ORDER'] = winnow_ordered_list(
fingerprint['SYNTACTIC']['HEADER_ORDER'])
return fingerprint
def test_terminate_join(self):
worker_pool = Pool(1, worker_names='WorkerThread')
worker_pool.terminate()
worker_pool.join()
def test_close_terminate(self):
worker_pool = Pool(1, worker_names='WorkerThread')
worker_pool.close()
worker_pool.terminate()