def setUp(self):
"""Set up."""
self.mock_compute = mock.patch.object(batch_enforcer.compute,
'ComputeClient').start()
self.gce_service = self.mock_compute().service
self.gce_service.networks().list().execute.return_value = (
constants.SAMPLE_TEST_NETWORK_SELFLINK)
self.project = constants.TEST_PROJECT
self.policy = json.loads(constants.RAW_EXPECTED_JSON_POLICY)
self.batch_enforcer = batch_enforcer.BatchFirewallEnforcer(
dry_run=True)
self.mock_time = mock.patch.object(batch_enforcer.datelib,
'Timestamp').start()
self.mock_time.now().AsMicroTimestamp.return_value = MOCK_TIMESTAMP
self.mock_time.now().AsSecondsSinceEpoch.return_value = MOCK_TIMESTAMP
self.expected_summary = (
enforcer_log_pb2.BatchResult(
batch_id=MOCK_TIMESTAMP,
timestamp_start_msec=MOCK_TIMESTAMP,
timestamp_end_msec=MOCK_TIMESTAMP))
self.expected_rules = copy.deepcopy(
constants.EXPECTED_FIREWALL_RULES.values())
self.addCleanup(mock.patch.stopall)
def initialize_batch_enforcer(concurrent_threads, max_write_threads,
max_running_operations, dry_run):
"""Initialize and return a BatchFirewallEnforcer object.
Args:
concurrent_threads: The number of parallel enforcement threads to execute.
max_write_threads: The maximum number of enforcement threads that can be
actively updating project firewalls.
max_running_operations: The maximum number of write operations per
enforcement thread.
dry_run: If True, will simply log what action would have been taken
without actually applying any modifications.
Returns:
A BatchFirewallEnforcer instance.
"""
if max_write_threads:
project_sema = threading.BoundedSemaphore(value=max_write_threads)
else:
project_sema = None
enforcer = batch_enforcer.BatchFirewallEnforcer(
dry_run=dry_run,
concurrent_workers=concurrent_threads,
project_sema=project_sema,
max_running_operations=max_running_operations)
return enforcer