use crate::observation::UsageObservation;
use crate::rating::RatedUsageRecord;
pub trait ObservationStore: Send + Sync {
fn append_observation(&self, observation: UsageObservation) -> Result<StoreResult, StoreError>;
}
pub trait RatedRecordStore: Send + Sync {
fn append_rated_record(&self, record: RatedUsageRecord) -> Result<StoreResult, StoreError>;
}
#[derive(Debug, Clone, PartialEq)]
pub enum StoreResult {
Appended,
AlreadyExists,
}
#[derive(Debug, Clone)]
pub enum StoreError {
ConnectionError(String),
IntegrityError(String),
SerializationError(String),
Other(String),
}
impl std::fmt::Display for StoreError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
StoreError::ConnectionError(e) => write!(f, "Connection error: {e}"),
StoreError::IntegrityError(e) => write!(f, "Integrity error: {e}"),
StoreError::SerializationError(e) => write!(f, "Serialization error: {e}"),
StoreError::Other(e) => write!(f, "Store error: {e}"),
}
}
}
impl std::error::Error for StoreError {}
#[cfg(test)]
mod tests {
use super::*;
use crate::identity::UsageEventId;
use crate::observation::{Attributes, MeterSet, UsageObservation, UsageOutcome, UsageSource, UsageTiming};
use crate::pricing::ModelRef;
use chrono::Utc;
use std::sync::{Arc, Mutex};
struct InMemoryObservationStore {
observations: Arc<Mutex<Vec<UsageObservation>>>,
}
impl InMemoryObservationStore {
fn new() -> Self {
Self {
observations: Arc::new(Mutex::new(Vec::new())),
}
}
}
impl ObservationStore for InMemoryObservationStore {
fn append_observation(&self, observation: UsageObservation) -> Result<StoreResult, StoreError> {
self.observations.lock().unwrap().push(observation);
Ok(StoreResult::Appended)
}
}
#[test]
fn in_memory_store_appends() {
let store = InMemoryObservationStore::new();
let obs = UsageObservation {
event_id: UsageEventId::from_raw("test-1"),
subject: crate::identity::BillingSubject::default(),
meter_set: MeterSet::new(),
model_ref: ModelRef {
billable_model: "test".to_string(),
vendor: None,
region: None,
tier: None,
},
provider_ref: None,
source: UsageSource::Estimated,
outcome: UsageOutcome::Success,
timing: UsageTiming {
observed_at: Utc::now(),
completed_at: None,
},
correlation: crate::identity::CorrelationIds::default(),
attributes: Attributes::new(),
};
let result = store.append_observation(obs).unwrap();
assert_eq!(result, StoreResult::Appended);
assert_eq!(store.observations.lock().unwrap().len(), 1);
}
#[test]
fn store_error_display() {
let err = StoreError::ConnectionError("db down".to_string());
assert_eq!(err.to_string(), "Connection error: db down");
}
#[test]
fn store_result_equality() {
assert_eq!(StoreResult::Appended, StoreResult::Appended);
assert_eq!(StoreResult::AlreadyExists, StoreResult::AlreadyExists);
assert_ne!(StoreResult::Appended, StoreResult::AlreadyExists);
}
}