Skip to content

Commit

Permalink
Fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
nabla-c0d3 committed Jun 28, 2020
1 parent e00b917 commit 93104bf
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 12 deletions.
12 changes: 6 additions & 6 deletions sslyze/scanner.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,7 @@ def __init__(
final_concurrent_server_scans_limit = concurrent_server_scans_limit
self._concurrent_server_scans_count = final_concurrent_server_scans_limit

self._all_thread_pools: List[ThreadPoolExecutor] = []
self._thread_pools: List[ThreadPoolExecutor] = []
self._queued_server_scans: List[_QueuedServerScan] = []

def _get_assigned_thread_pool_index(self) -> int:
Expand All @@ -172,9 +172,9 @@ def _get_assigned_thread_pool_index(self) -> int:
assigned_thread_pool_index = currently_queued_scans_count % allowed_thread_pools_count

try:
self._all_thread_pools[assigned_thread_pool_index]
self._thread_pools[assigned_thread_pool_index]
except IndexError:
self._all_thread_pools.append(ThreadPoolExecutor(max_workers=self._per_server_concurrent_connections_count))
self._thread_pools.append(ThreadPoolExecutor(max_workers=self._per_server_concurrent_connections_count))

return assigned_thread_pool_index

Expand All @@ -190,7 +190,7 @@ def queue_scan(self, server_scan: ServerScanRequest) -> None:

# Assign the server to scan to a thread pool
assigned_thread_pool_index = self._get_assigned_thread_pool_index()
assigned_thread_pool = self._all_thread_pools[assigned_thread_pool_index]
assigned_thread_pool = self._thread_pools[assigned_thread_pool_index]

# Convert each scan command within the server scan request into jobs
queued_futures_per_scan_command: Dict[ScanCommandType, Set[Future]] = {}
Expand Down Expand Up @@ -301,9 +301,9 @@ def get_results(self) -> Iterable[ServerScanResult]:

def _shutdown_thread_pools(self) -> None:
self._queued_server_scans = []
for thread_pool in self._all_thread_pools:
for thread_pool in self._thread_pools:
thread_pool.shutdown(wait=True)
self._all_thread_pools = []
self._thread_pools = []

# Force garbage collection because for some reason the Future objects created by ThreadPoolExecutor.submit()
# take a ton of memory (compared to what they do - holding a function to call and its arguments):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,7 @@ def test_certificate_with_no_cn(self):
# When running the scan, it succeeds
plugin_result = CertificateInfoImplementation.scan_server(server_info)

assert plugin_result.certificate_deployments[0].verified_certificate_chain
assert plugin_result.certificate_deployments[0].received_certificate_chain

def test_certificate_with_no_subject(self):
# Given a server to scan that has a certificate with no Subject
Expand Down
13 changes: 8 additions & 5 deletions tests/test_scanner.py
Original file line number Diff line number Diff line change
Expand Up @@ -232,16 +232,16 @@ def test(self):

# And the right number of scans was performed
assert total_server_scans_count == len(scanner._queued_server_scans)
assert total_server_scans_count == len(scanner._server_to_thread_pool)

# And the chosen network settings were used
assert concurrent_server_scans_limit == len(scanner._all_thread_pools)
for pool in scanner._all_thread_pools:
assert concurrent_server_scans_limit == len(scanner._thread_pools)
for pool in scanner._thread_pools:
assert per_server_concurrent_connections_limit == pool._max_workers

# And the server scans were evenly distributed among the thread pools to maximize performance
expected_server_scans_per_pool = int(total_server_scans_count / concurrent_server_scans_limit)
server_scans_per_pool_count = Counter(scanner._server_to_thread_pool.values())
thread_pools_used = [server_scan.queued_on_thread_pool_at_index for server_scan in scanner._queued_server_scans]
server_scans_per_pool_count = Counter(thread_pools_used)
for pool_count in server_scans_per_pool_count.values():
assert expected_server_scans_per_pool == pool_count

Expand All @@ -266,5 +266,8 @@ def test_emergency_shutdown(self):
scanner.emergency_shutdown()

# And all the queued jobs were done or cancelled
for completed_future in as_completed(scanner._queued_future_to_server_and_scan_cmd.keys()):
all_queued_futures = []
for server_scan in scanner._queued_server_scans:
all_queued_futures.extend(server_scan.all_queued_scan_jobs)
for completed_future in as_completed(all_queued_futures):
assert completed_future.done()

0 comments on commit 93104bf

Please sign in to comment.