use polyc_proto::proto::polychrome::events::v1::UsageEvent;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct UsageFact {
pub input_tokens: u64,
pub output_tokens: u64,
}
impl From<UsageEvent> for UsageFact {
fn from(event: UsageEvent) -> Self {
Self {
input_tokens: event.input_tokens,
output_tokens: event.output_tokens,
}
}
}
pub fn fold_usage_event(payload: &[u8]) -> Result<UsageFact, buffa::DecodeError> {
polyc_proto::events_decode::try_decode_event_payload::<UsageEvent>(payload).map(UsageFact::from)
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs, clippy::unwrap_used)]
use buffa::Message as _;
use super::*;
#[test]
fn round_trips_a_usage_payload() {
let event = UsageEvent {
input_tokens: 9,
output_tokens: 3,
__buffa_unknown_fields: buffa::UnknownFields::default(),
};
let bytes = event.encode_to_vec();
let fact = fold_usage_event(&bytes).expect("decode usage");
assert_eq!(fact.input_tokens, 9);
assert_eq!(fact.output_tokens, 3);
}
#[test]
fn empty_payload_decodes_to_zero_zero() {
let fact = fold_usage_event(&[]).expect("empty payload decodes");
assert_eq!(fact, UsageFact::default());
}
#[test]
fn garbage_bytes_return_an_error() {
assert!(fold_usage_event(&[0xFF, 0xFE, 0xFD]).is_err());
}
}