def assert_limits(self, _, patch_set_cpu, patch_set_memory_limit, patch_get_enforce, patch_add_event, ext_name, expected_cpu_limit, limits_enforced=True, exception_raised=False): should_limit = expected_cpu_limit > 0 patch_get_enforce.return_value = limits_enforced if exception_raised: patch_set_memory_limit.side_effect = CGroupsException( 'set_memory_limit error') try: cg = CGroupConfigurator.for_extension(ext_name) cg.set_limits() if exception_raised: self.fail('exception expected') except CGroupsException: if not exception_raised: self.fail('exception not expected') self.assertEqual(should_limit, patch_set_cpu.called) self.assertEqual(should_limit, patch_set_memory_limit.called) self.assertEqual(should_limit, patch_add_event.called) if should_limit: actual_cpu_limit = patch_set_cpu.call_args[0][0] actual_memory_limit = patch_set_memory_limit.call_args[0][0] event_kw_args = patch_add_event.call_args[1] self.assertEqual(expected_cpu_limit, actual_cpu_limit) self.assertTrue(actual_memory_limit >= DEFAULT_MEM_LIMIT_MIN_MB) self.assertEqual(event_kw_args['op'], 'SetCGroupsLimits') self.assertEqual(event_kw_args['is_success'], not exception_raised) self.assertTrue( '{0}%'.format(expected_cpu_limit) in event_kw_args['message']) self.assertTrue(ext_name in event_kw_args['message']) self.assertEqual( exception_raised, 'set_memory_limit error' in event_kw_args['message'])
def test_telemetry_instantiation_as_superuser(self): """ Tracking a new cgroup for an extension; collect all metrics. """ # Record initial state initial_cgroup = make_self_cgroups() # Put the process into a different cgroup, consume some resources, ensure we see them end-to-end test_cgroup = CGroupConfigurator.for_extension("agent_unittest") test_cgroup.add(os.getpid()) self.assertNotEqual(initial_cgroup.cgroups['cpu'], test_cgroup.cgroups['cpu']) self.assertNotEqual(initial_cgroup.cgroups['memory'], test_cgroup.cgroups['memory']) self.exercise_telemetry_instantiation(test_cgroup) # Restore initial state CGroupsTelemetry.stop_tracking("agent_unittest") initial_cgroup.add(os.getpid())