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 FetchClientResourceUsageRecords(
client_id: str, start_range: timestamp_pb2.Timestamp,
end_range: timestamp_pb2.Timestamp
) -> List[resource_pb2.ClientResourceUsageRecord]:
"""Returns aggregated resource usage metrics of a client from Fleetspeak.
Args:
client_id: Id of the client to fetch Fleetspeak resource usage records for.
start_range: Start timestamp of range.
end_range: end timestamp of range.
Returns:
A list of client resource usage records retrieved from Fleetspeak.
"""
res = fleetspeak_connector.CONN.outgoing.FetchClientResourceUsageRecords(
admin_pb2.FetchClientResourceUsageRecordsRequest(
client_id=GRRIDToFleetspeakID(client_id),
start_timestamp=start_range,
end_timestamp=end_range))
if not res.records:
return []
return list(res.records)