def testFetchClientResourceUsageRecords(self):
conn = mock.MagicMock()
conn.outgoing.FetchClientResourceUsageRecords.return_value = admin_pb2.FetchClientResourceUsageRecordsResponse(
records=[{
"mean_user_cpu_rate": 1,
"max_system_cpu_rate": 2
}, {
"mean_user_cpu_rate": 4,
"max_system_cpu_rate": 8
}])
with mock.patch.object(fleetspeak_connector, "CONN", conn):
expected_records_list = [
resource_pb2.ClientResourceUsageRecord(
mean_user_cpu_rate=1, max_system_cpu_rate=2),
resource_pb2.ClientResourceUsageRecord(
mean_user_cpu_rate=4, max_system_cpu_rate=8)
]
arbitrary_date = timestamp_pb2.Timestamp(seconds=1, nanos=1)
self.assertListEqual(
fleetspeak_utils.FetchClientResourceUsageRecords(
client_id=_TEST_CLIENT_ID,
start_range=arbitrary_date,
end_range=arbitrary_date), expected_records_list)
conn.outgoing.FetchClientResourceUsageRecords.assert_called_once()
args, _ = conn.outgoing.FetchClientResourceUsageRecords.call_args
self.assertEqual(
args[0],
admin_pb2.FetchClientResourceUsageRecordsRequest(
client_id=fleetspeak_utils.GRRIDToFleetspeakID(_TEST_CLIENT_ID),
start_timestamp=arbitrary_date,
end_timestamp=arbitrary_date))
def _MockConnReturningRecords(client_ruds):
conn = mock.MagicMock()
records = []
for record in client_ruds:
records.append(
resource_pb2.ClientResourceUsageRecord(
scope=record["scope"],
pid=record["pid"],
process_start_time=record["process_start_time"],
client_timestamp=record["client_timestamp"],
server_timestamp=record["server_timestamp"],
mean_user_cpu_rate=record["mean_user_cpu_rate"],
max_user_cpu_rate=record["max_user_cpu_rate"],
mean_system_cpu_rate=record["mean_system_cpu_rate"],
max_system_cpu_rate=record["max_system_cpu_rate"],
mean_resident_memory_mib=record["mean_resident_memory_mib"],
max_resident_memory_mib=record["max_resident_memory_mib"]))
conn.outgoing.FetchClientResourceUsageRecords.return_value = admin_pb2.FetchClientResourceUsageRecordsResponse(
records=records)
return conn
