github-ops
76 lines
1.9 KiB
Rust
76 lines
1.9 KiB
Rust
use std::sync::Arc;
|
|
|
|
use async_trait::async_trait;
|
|
|
|
use crate::{TenantId, UsageRollup};
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub enum BillingGate {
|
|
Allow,
|
|
Deny { reason: String },
|
|
}
|
|
|
|
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
|
|
pub enum BillingError {
|
|
#[error("billing integration is not available in this edition")]
|
|
NotSupportedInEdition,
|
|
#[error("billing provider failure: {0}")]
|
|
Provider(String),
|
|
}
|
|
|
|
#[async_trait]
|
|
pub trait BillingHook: Send + Sync {
|
|
async fn on_rollup(&self, rollup: UsageRollup) -> Result<(), BillingError>;
|
|
|
|
async fn check_billing_gate(&self, tenant: &TenantId) -> Result<BillingGate, BillingError>;
|
|
}
|
|
|
|
pub type SharedBillingHook = Arc<dyn BillingHook>;
|
|
|
|
#[derive(Clone, Copy, Debug, Default)]
|
|
pub struct NoopBillingHook;
|
|
|
|
#[async_trait]
|
|
impl BillingHook for NoopBillingHook {
|
|
async fn on_rollup(&self, _rollup: UsageRollup) -> Result<(), BillingError> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn check_billing_gate(&self, _tenant: &TenantId) -> Result<BillingGate, BillingError> {
|
|
Ok(BillingGate::Allow)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{BillingGate, BillingHook, NoopBillingHook};
|
|
use crate::{TenantId, UsagePeriod, UsageRollup, WorkspaceId};
|
|
|
|
#[tokio::test]
|
|
async fn noop_billing_hook_allows_default_gate_and_rollups() {
|
|
let hook = NoopBillingHook;
|
|
|
|
hook.on_rollup(UsageRollup {
|
|
workspace_id: WorkspaceId::new("ws_01"),
|
|
agent_id: None,
|
|
operation_id: None,
|
|
period: UsagePeriod::Last24Hours,
|
|
calls_total: 10,
|
|
calls_ok: 9,
|
|
calls_error: 1,
|
|
p50_ms: 12,
|
|
p95_ms: 42,
|
|
p99_ms: 77,
|
|
})
|
|
.await
|
|
.unwrap();
|
|
|
|
let gate = hook
|
|
.check_billing_gate(&TenantId::new("tenant_01"))
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(gate, BillingGate::Allow);
|
|
}
|
|
}
|