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()