use std::time::Duration;
use distributed::{
sourced, Aggregate, AggregateBuilder, AsyncOutboxStore, ClaimOutboxMessages, CommitBatch,
Entity, GetStream, HashMapRepository, InMemorySnapshotStore, OutboxMessage, ReadModel,
ReadModelWritePlanBuilder, RelationalReadModel, RelationalReadModelQueryStore, RepositoryError,
RowKey, RowValue, SnapshotRecord, SnapshotStore, Snapshottable, StreamIdentity, StreamWrite,
TransactionalCommit, Versioned,
};
use serde::{Deserialize, Serialize};
#[derive(Default)]
struct AlphaAggregate {
entity: Entity,
}
#[sourced(entity, aggregate_type = "async.alpha")]
impl AlphaAggregate {
#[event("touched")]
fn touch(&mut self, id: String) {
self.entity.set_id(&id);
}
}
#[derive(Default)]
struct BetaAggregate {
entity: Entity,
}
#[sourced(entity, aggregate_type = "async.beta")]
impl BetaAggregate {
#[event("touched")]
fn touch(&mut self, id: String) {
self.entity.set_id(&id);
}
}
#[derive(Default)]
struct SnapshotCounter {
entity: Entity,
value: i32,
}
#[sourced(entity, aggregate_type = "async.snapshot_counter")]
impl SnapshotCounter {
#[event("incremented")]
fn increment(&mut self, id: String, by: i32) {
self.entity.set_id(&id);
self.value += by;
}
}
impl Snapshottable for SnapshotCounter {
type Snapshot = i32;
fn create_snapshot(&self) -> Self::Snapshot {
self.value
}
fn restore_from_snapshot(&mut self, snapshot: Self::Snapshot) {
self.value = snapshot;
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, ReadModel)]
#[table("repository_test_views")]
struct TestView {
#[id]
id: String,
value: i32,
}
fn test_view_key(id: &str) -> RowKey {
RowKey::new([("id", RowValue::String(id.into()))])
}
async fn load_test_view(repo: &HashMapRepository, id: &str) -> Option<Versioned<TestView>> {
let request = ReadModelWritePlanBuilder::new()
.load::<TestView>(test_view_key(id))
.unwrap();
let graph = repo.load_graph(request).await.unwrap();
graph.root.map(|root| Versioned {
data: TestView::from_row(root.data).unwrap(),
version: root.version,
})
}
#[tokio::test]
async fn aggregate_repository_separates_streams_by_aggregate_type() {
let repo = HashMapRepository::new();
let alpha_repo = repo.clone().aggregate::<AlphaAggregate>();
let beta_repo = repo.clone().aggregate::<BetaAggregate>();
let mut alpha = AlphaAggregate::default();
alpha.touch("shared-id".into()).unwrap();
let mut beta = BetaAggregate::default();
beta.touch("shared-id".into()).unwrap();
alpha_repo.commit(&mut alpha).await.unwrap();
beta_repo.commit(&mut beta).await.unwrap();
let loaded_alpha = alpha_repo.get("shared-id").await.unwrap().unwrap();
let loaded_beta = beta_repo.get("shared-id").await.unwrap().unwrap();
assert_eq!(loaded_alpha.entity().events().len(), 1);
assert_eq!(loaded_beta.entity().events().len(), 1);
}
#[tokio::test]
async fn batch_rejects_duplicate_stream_identity_before_write() {
let repo = HashMapRepository::new();
let identity = StreamIdentity::new("async.alpha", "duplicate").unwrap();
let mut first = Entity::with_id("duplicate");
first.digest_empty("first_recorded").unwrap();
let mut second = Entity::with_id("duplicate");
second.digest_empty("second_recorded").unwrap();
let err = repo
.commit_batch(CommitBatch::new(vec![
StreamWrite::new(identity.clone(), &mut first),
StreamWrite::new(identity.clone(), &mut second),
]))
.await
.unwrap_err();
assert!(
matches!(
&err,
RepositoryError::DuplicateStreamInBatch { id } if id == "async.alpha:duplicate"
),
"unexpected error: {err}"
);
assert!(repo.get_stream(&identity).await.unwrap().is_none());
}
#[tokio::test]
async fn read_model_write_plan_can_commit_against_store() {
let repo = HashMapRepository::new();
let view = TestView {
id: "view-1".into(),
value: 42,
};
let mut session = ReadModelWritePlanBuilder::new();
session.upsert(&view).unwrap();
let outcome = session.commit(&repo).await.unwrap();
let loaded = load_test_view(&repo, "view-1").await.unwrap();
assert!(outcome.was_applied());
assert_eq!(loaded.data, view);
}
#[tokio::test]
async fn snapshot_store_uses_full_stream_identity() {
let store = InMemorySnapshotStore::new();
let alpha = StreamIdentity::new("async.alpha", "same-id").unwrap();
let beta = StreamIdentity::new("async.beta", "same-id").unwrap();
store
.save_snapshot(
&alpha,
SnapshotRecord::new("async.alpha", "same-id", 1, 1, vec![1]),
)
.await
.unwrap();
store
.save_snapshot(
&beta,
SnapshotRecord::new("async.beta", "same-id", 2, 1, vec![2]),
)
.await
.unwrap();
let loaded_alpha = store.get_snapshot(&alpha).await.unwrap().unwrap();
let loaded_beta = store.get_snapshot(&beta).await.unwrap().unwrap();
assert_eq!(loaded_alpha.version, 1);
assert_eq!(loaded_beta.version, 2);
assert_eq!(loaded_alpha.aggregate_type, "async.alpha");
assert_eq!(loaded_beta.aggregate_type, "async.beta");
}
#[tokio::test]
async fn snapshot_repository_writes_cache_without_event_record() {
let repo = HashMapRepository::new();
let snapshot_repo = repo
.clone()
.aggregate::<SnapshotCounter>()
.with_snapshots(2);
let id = "snapshot-counter-1";
let mut counter = SnapshotCounter::default();
counter.increment(id.into(), 2).unwrap();
snapshot_repo.commit(&mut counter).await.unwrap();
counter.increment(id.into(), 3).unwrap();
snapshot_repo.commit(&mut counter).await.unwrap();
let identity = StreamIdentity::new(SnapshotCounter::aggregate_type(), id).unwrap();
let stream = repo.get_stream(&identity).await.unwrap().unwrap();
let snapshot = repo.get_snapshot(&identity).await.unwrap().unwrap();
assert_eq!(stream.events().len(), 2);
assert_eq!(stream.events()[0].event_name, "incremented");
assert_eq!(stream.events()[1].event_name, "incremented");
assert_eq!(snapshot.version, 2);
assert_eq!(snapshot.aggregate_type, SnapshotCounter::aggregate_type());
assert_eq!(snapshot.payload, bitcode::serialize(&5_i32).unwrap());
}
#[tokio::test]
async fn snapshot_repository_ignores_invalid_cache_and_replays_events() {
let repo = HashMapRepository::new();
let aggregate_repo = repo.clone().aggregate::<SnapshotCounter>();
let snapshot_repo = repo
.clone()
.aggregate::<SnapshotCounter>()
.with_snapshots(10);
let id = "snapshot-counter-invalid";
let mut counter = SnapshotCounter::default();
counter.increment(id.into(), 4).unwrap();
counter.increment(id.into(), 6).unwrap();
aggregate_repo.commit(&mut counter).await.unwrap();
let identity = StreamIdentity::new(SnapshotCounter::aggregate_type(), id).unwrap();
let mut invalid = SnapshotRecord::new(SnapshotCounter::aggregate_type(), id, 1, 1, vec![0xff]);
invalid.payload_codec = "json".into();
repo.save_snapshot(&identity, invalid).await.unwrap();
let loaded = snapshot_repo.get(id).await.unwrap().unwrap();
assert_eq!(loaded.value, 10);
assert_eq!(loaded.entity().snapshot_version(), 0);
}
#[tokio::test]
async fn snapshot_repository_ignores_cache_past_stream_version_and_replays_events() {
let repo = HashMapRepository::new();
let aggregate_repo = repo.clone().aggregate::<SnapshotCounter>();
let snapshot_repo = repo
.clone()
.aggregate::<SnapshotCounter>()
.with_snapshots(10);
let id = "snapshot-counter-ahead";
let mut counter = SnapshotCounter::default();
counter.increment(id.into(), 4).unwrap();
aggregate_repo.commit(&mut counter).await.unwrap();
let identity = StreamIdentity::new(SnapshotCounter::aggregate_type(), id).unwrap();
let record = SnapshotRecord::new(
SnapshotCounter::aggregate_type(),
id,
2,
1,
bitcode::serialize(&999_i32).unwrap(),
);
repo.save_snapshot(&identity, record).await.unwrap();
let loaded = snapshot_repo.get(id).await.unwrap().unwrap();
assert_eq!(loaded.value, 4);
assert_eq!(loaded.entity().version(), 1);
assert_eq!(loaded.entity().snapshot_version(), 0);
}
#[tokio::test]
async fn outbox_repository_delegates_worker_operations() {
let repo = HashMapRepository::new();
let outbox = repo.outbox_store();
let message = OutboxMessage::create("msg-1", "alpha.happened", b"{}".to_vec()).unwrap();
let mut aggregate = AlphaAggregate::default();
aggregate.touch("outbox-aggregate-1".into()).unwrap();
repo.clone()
.aggregate::<AlphaAggregate>()
.outbox(message)
.commit(&mut aggregate)
.await
.unwrap();
let claimed = outbox
.claim_async(ClaimOutboxMessages::new(
"worker-1",
1,
Duration::from_secs(60),
))
.await
.unwrap();
assert_eq!(claimed.len(), 1);
assert_eq!(claimed[0].worker_id.as_deref(), Some("worker-1"));
}