def test_copy_all_binaries(self, mock_logger: mock.MagicMock):
"""Test the top-level copy function with real multiprocessing.
Note that coverage doesn't see code run by other processes, so it doesn't show any coverage
for the CopyTask or Consumer classes, but we do in fact run all of the code.
"""
mock_main = MockMain()
from lambda_functions.downloader import copy_all
copy_all.main = mock_main
copy_all.copy_all_binaries()
# Verify that every binary "in CarbonBlack" was sent to a Consumer process.
self.assertEqual(NUM_BINARIES, mock_main.index.value) # type: ignore
md5s = [mock_main.download_invocations[i] for i in range(NUM_BINARIES)]
self.assertEqual(list(range(NUM_BINARIES)), sorted(md5s))
# Verify that a logger was created for the main process and for each consumer.
# Also verify each log statement from the main logger.
mock_logger.assert_has_calls(
[mock.call('carbon_black_copy')] +
[mock.call().debug('Enqueuing %s', mock.ANY)] * NUM_BINARIES +
[mock.call().info('All CopyTasks Finished!')],
any_order=True
)
def test_copy_with_errors(self, mock_logger: mock.MagicMock):
"""Test top-level copy function with injected errors."""
mock_main = MockMain(inject_errors=True)
from lambda_functions.downloader import copy_all
importlib.reload(copy_all) # Reloading is required for the mock logger to be reset.
copy_all.main = mock_main
copy_all.copy_all_binaries()
# Verify that the root logger logged all of the failed binaries.
mock_logger.assert_has_calls([
mock.call().error('%d %s failed to copy: \n%s', mock.ANY, 'binaries', mock.ANY)
])
def cb_copy_all(self) -> None:
"""Copy all binaries from CarbonBlack into BinaryAlert.
Raises:
InvalidConfigError: If the CarbonBlack downloader is not enabled.
"""
if not self._config.enable_carbon_black_downloader:
raise InvalidConfigError('CarbonBlack downloader is not enabled.')
os.environ['CARBON_BLACK_URL'] = self._config.carbon_black_url
os.environ['ENCRYPTED_CARBON_BLACK_API_TOKEN'] = (
self._config.encrypted_carbon_black_api_token
)
os.environ['TARGET_S3_BUCKET'] = self._config.binaryalert_s3_bucket_name
# Downloader must be imported here because the cb_api is configured at import time.
from lambda_functions.downloader import copy_all
copy_all.copy_all_binaries()