macltok

#![cfg(feature = "mirror")]

use std::time::Duration;

use futures_util::StreamExt;

use macltok::{

    ClientHandle, Command, ConnectionState, Engine, EngineConfig, Event, MentionCoverage,

    MentionsBoundary, MentionsPage, MentionsUpdate, Mirror, MirrorHandle, StateChange,

    Subscription, TimelineChange, TimelineCommand, TimelineDiff, TimelineId, TimelineItem, Update,

    UploadOutcome,

};

use matrix_sdk::{

    ruma::{

        OwnedRoomId, OwnedTransactionId, TransactionId, api::client::sync::sync_events::v5,

        event_id, room_id, user_id,

    },

    test_utils::mocks::MatrixMockServer,

};

use tokio::sync::mpsc;

struct Harness {

    handle: ClientHandle,

    mirror: MirrorHandle,

    inject: mpsc::UnboundedSender<Event>,

    room_id: OwnedRoomId,

    _server: MatrixMockServer,

}

/// the mirror reads the real engine's stream with synthetic events mixed in, so a test can

/// hand it a diff the engine never would

async fn start(config: EngineConfig, capacity: usize) -> Harness {

    let server = MatrixMockServer::new().await;

    server

        .mock_sliding_sync()

        .ok(v5::Response::new("1".to_owned()))

        .mount()

        .await;

    let client = server.client_builder().build().await;

    let room_id = room_id!("!mirror:localhost");

    server.mock_room_state_encryption().plain().mount().await;

    server.sync_joined_room(&client, room_id).await;

    let (engine, handle, events) =

        Engine::from_client(client, config.with_initial_history_limit(0))

            .await

            .unwrap();

    tokio::spawn(engine.run());

    let (inject, mut injected) = mpsc::unbounded_channel();

    let injected = futures_util::stream::poll_fn(move |context| injected.poll_recv(context));

    let (mirror, mirror_handle) = Mirror::with_capacity(

        handle.clone(),

        futures_util::stream::select(events, injected),

        capacity,

    );

    tokio::spawn(mirror.run());

    Harness {

        handle,

        mirror: mirror_handle,

        inject,

        room_id: room_id.to_owned(),

        _server: server,

    }

}

async fn next_matching(

    subscription: &mut Subscription,

    predicate: impl Fn(&StateChange) -> bool,

) -> Update {

    tokio::time::timeout(Duration::from_secs(5), async {

        loop {

            let update = subscription.next().await.expect("the mirror stopped");

            if predicate(&update.change) {

                return update;

            }

        }

    })

    .await

    .expect("expected update did not arrive")

}

fn timeline_change(

    change: &StateChange,

    id: &TimelineId,

    predicate: impl Fn(&TimelineChange) -> bool,

) -> bool {

    matches!(change, StateChange::Timeline { timeline_id, change } if timeline_id == id && predicate(change))

}

fn item(timestamp_ms: u64) -> TimelineItem {

    TimelineItem::DayDivider { timestamp_ms }

}

fn reset(timeline_id: &TimelineId, items: Vec<TimelineItem>) -> Event {

    Event::TimelineReset {

        timeline_id: timeline_id.clone(),

        items: items.into(),

    }

}

/// the room never opens, so every resync is refused and the divergence holds still to be looked at

#[tokio::test]

async fn diverged_batch_keeps_last_good_items_until_a_reset() {

    let harness = start(EngineConfig::default(), 256).await;

    let timeline_id = TimelineId::from(harness.room_id.clone());

    let mut updates = harness.mirror.subscribe();

    harness

        .inject

        .send(reset(&timeline_id, vec![item(1), item(2)]))

        .unwrap();

    for diffs in [

        vec![

            TimelineDiff::PushBack { value: item(3) },

            TimelineDiff::Remove { index: 7 },

        ],

        vec![TimelineDiff::PushBack { value: item(4) }],

    ] {

        harness

            .inject

            .send(Event::TimelineChanged {

                timeline_id: timeline_id.clone(),

                diffs,

            })

            .unwrap();

    }

    next_matching(&mut updates, |change| {

        timeline_change(change, &timeline_id, |change| {

            matches!(change, TimelineChange::Diverged)

        })

    })

    .await;

    let snapshot = harness.mirror.snapshot();

    let timeline = &snapshot.timelines[&timeline_id];

    assert!(timeline.diverged);

    assert_eq!(timeline.items, [item(1), item(2)].into_iter().collect());

    harness

        .inject

        .send(reset(&timeline_id, vec![item(5)]))

        .unwrap();

    let update = next_matching(&mut updates, |change| {

        timeline_change(change, &timeline_id, |_| true)

    })

    .await;

    assert!(

        timeline_change(&update.change, &timeline_id, |change| {

            matches!(change, TimelineChange::Reset { .. })

        }),

        "a diverged timeline took {update:?}"

    );

    let snapshot = harness.mirror.snapshot();

    let timeline = &snapshot.timelines[&timeline_id];

    assert!(!timeline.diverged);

    assert_eq!(timeline.items, [item(5)].into_iter().collect());

}

#[tokio::test]

async fn a_diff_for_a_timeline_without_a_reset_diverges_it() {

    let harness = start(EngineConfig::default(), 256).await;

    let timeline_id = TimelineId::from(harness.room_id.clone());

    let mut updates = harness.mirror.subscribe();

    harness

        .inject

        .send(Event::TimelineChanged {

            timeline_id: timeline_id.clone(),

            diffs: vec![TimelineDiff::PushBack { value: item(1) }],

        })

        .unwrap();

    next_matching(&mut updates, |change| {

        timeline_change(change, &timeline_id, |change| {

            matches!(change, TimelineChange::Diverged)

        })

    })

    .await;

    let snapshot = harness.mirror.snapshot();

    let timeline = &snapshot.timelines[&timeline_id];

    assert!(timeline.diverged);

    assert!(timeline.items.is_empty());

}

#[tokio::test]

async fn room_and_thread_timelines_stay_apart() {

    let harness = start(EngineConfig::default(), 256).await;

    let room = TimelineId::from(harness.room_id.clone());

    let thread = TimelineId::thread(

        harness.room_id.clone(),

        event_id!("$root:localhost").to_owned(),

    );

    let mut updates = harness.mirror.subscribe();

    harness.inject.send(reset(&room, vec![item(1)])).unwrap();

    harness.inject.send(reset(&thread, vec![item(2)])).unwrap();

    harness

        .inject

        .send(Event::TimelineChanged {

            timeline_id: thread.clone(),

            diffs: vec![TimelineDiff::PushBack { value: item(3) }],

        })

        .unwrap();

    next_matching(&mut updates, |change| {

        timeline_change(change, &thread, |change| {

            matches!(change, TimelineChange::Changed { .. })

        })

    })

    .await;

    let snapshot = harness.mirror.snapshot();

    assert_eq!(

        snapshot.timelines[&room].items,

        [item(1)].into_iter().collect()

    );

    assert_eq!(

        snapshot.timelines[&thread].items,

        [item(2), item(3)].into_iter().collect()

    );

}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]

async fn subscribers_joining_mid_stream_see_no_gap_and_nothing_twice() {

    const CHUNKS: u64 = 64;

    const CHUNK: u64 = 32;

    let harness = start(EngineConfig::default(), 4096).await;

    let transaction_id = TransactionId::new();

    let mut followers = Vec::new();

    for chunk in 0..CHUNKS {

        for uploaded_bytes in chunk * CHUNK + 1..=(chunk + 1) * CHUNK {

            harness

                .inject

                .send(Event::MediaUploadProgress {

                    room_id: harness.room_id.clone(),

                    transaction_id: transaction_id.clone(),

                    uploaded_bytes,

                    total_bytes: CHUNKS * CHUNK,

                    complete: false,

                })

                .unwrap();

        }

        followers.push(tokio::spawn(follow(

            harness.mirror.subscribe(),

            transaction_id.clone(),

            CHUNKS * CHUNK,

        )));

    }

    for follower in followers {

        follower.await.unwrap();

    }

}

async fn follow(mut updates: Subscription, transaction_id: OwnedTransactionId, total: u64) {

    let first = updates.next().await.unwrap();

    let StateChange::State(state) = &first.change else {

        panic!("a subscription opened with {first:?}");

    };

    let mut seen = state

        .uploads

        .get(&transaction_id)

        .map_or(0, |upload| upload.uploaded_bytes);

    let mut sequence = first.sequence;

    while seen < total {

        let update = next_matching(&mut updates, |_| true).await;

        assert_eq!(update.sequence, sequence + 1);

        sequence = update.sequence;

        if let StateChange::Upload { upload, .. } = update.change {

            assert_eq!(upload.uploaded_bytes, seen + 1);

            seen += 1;

        }

    }

}

#[tokio::test]

async fn lagging_subscriber_starts_over_from_a_fresh_state() {

    let harness = start(EngineConfig::default(), 4).await;

    let transaction_id = TransactionId::new();

    let mut updates = harness.mirror.subscribe();

    assert!(matches!(

        updates.next().await.unwrap().change,

        StateChange::State(_)

    ));

    for uploaded_bytes in 1..=20 {

        harness

            .inject

            .send(Event::MediaUploadProgress {

                room_id: harness.room_id.clone(),

                transaction_id: transaction_id.clone(),

                uploaded_bytes,

                total_bytes: 20,

                complete: uploaded_bytes == 20,

            })

            .unwrap();

    }

    tokio::time::timeout(Duration::from_secs(5), async {

        while !harness

            .mirror

            .snapshot()

            .uploads

            .get(&transaction_id)

            .is_some_and(|upload| upload.complete)

        {

            tokio::task::yield_now().await;

        }

    })

    .await

    .expect("the mirror never caught up");

    let update = updates.next().await.unwrap();

    let StateChange::State(state) = &update.change else {

        panic!("a lagged subscription carried on with {update:?}");

    };

    assert!(state.uploads[&transaction_id].complete);

}

#[tokio::test]

async fn resync_does_not_deadlock_on_full_queues() {

    let config = EngineConfig::default()

        .with_command_capacity(1)

        .with_event_capacity(1);

    let harness = start(config, 1024).await;

    let timeline_id = TimelineId::from(harness.room_id.clone());

    let mut updates = harness.mirror.subscribe();

    harness

        .handle

        .open_room(harness.room_id.clone())

        .await

        .unwrap();

    next_matching(&mut updates, |change| {

        timeline_change(change, &timeline_id, |change| {

            matches!(change, TimelineChange::Reset { .. })

        })

    })

    .await;

    let flood = tokio::spawn({

        let handle = harness.handle.clone();

        let timeline_id = timeline_id.clone();

        async move {

            loop {

                let command = Command::Timeline(TimelineCommand::Paginate {

                    timeline_id: timeline_id.clone(),

                    limit: 1,

                });

                if handle.send(command).await.is_err() {

                    return;

                }

            }

        }

    });

    for _ in 0..10 {

        harness

            .inject

            .send(Event::TimelineChanged {

                timeline_id: timeline_id.clone(),

                diffs: vec![TimelineDiff::Remove { index: 1000 }],

            })

            .unwrap();

        next_matching(&mut updates, |change| {

            timeline_change(change, &timeline_id, |change| {

                matches!(change, TimelineChange::Diverged)

            })

        })

        .await;

        next_matching(&mut updates, |change| {

            timeline_change(change, &timeline_id, |change| {

                matches!(change, TimelineChange::Reset { .. })

            })

        })

        .await;

    }

    flood.abort();

}

#[tokio::test]

async fn closing_a_timeline_evicts_it() {

    let harness = start(EngineConfig::default(), 256).await;

    let timeline_id = TimelineId::from(harness.room_id.clone());

    let mut updates = harness.mirror.subscribe();

    harness

        .handle

        .open_room(harness.room_id.clone())

        .await

        .unwrap();

    harness

        .handle

        .close_room(harness.room_id.clone())

        .await

        .unwrap();

    next_matching(&mut updates, |change| {

        timeline_change(change, &timeline_id, |change| {

            matches!(change, TimelineChange::Closed)

        })

    })

    .await;

    assert!(

        !harness

            .mirror

            .snapshot()

            .timelines

            .contains_key(&timeline_id)

    );

}

#[tokio::test]

async fn an_upload_ends_on_finished() {

    let harness = start(EngineConfig::default(), 256).await;

    let transaction_id = TransactionId::new();

    let mut updates = harness.mirror.subscribe();

    harness

        .inject

        .send(Event::MediaUploadProgress {

            room_id: harness.room_id.clone(),

            transaction_id: transaction_id.clone(),

            uploaded_bytes: 5,

            total_bytes: 10,

            complete: false,

        })

        .unwrap();

    next_matching(&mut updates, |change| {

        matches!(change, StateChange::Upload { .. })

    })

    .await;

    assert!(

        harness

            .mirror

            .snapshot()

            .uploads

            .contains_key(&transaction_id)

    );

    harness

        .inject

        .send(Event::MediaUploadFinished {

            room_id: harness.room_id.clone(),

            transaction_id: transaction_id.clone(),

            outcome: UploadOutcome::Cancelled,

        })

        .unwrap();

    next_matching(&mut updates, |change| {

        matches!(change, StateChange::UploadFinished { .. })

    })

    .await;

    assert!(harness.mirror.snapshot().uploads.is_empty());

}

#[tokio::test]

async fn mentions_updates_require_the_current_account_and_revision() {

    let harness = start(EngineConfig::default(), 256).await;

    let mut updates = harness.mirror.subscribe();

    let account = tokio::time::timeout(Duration::from_secs(5), async {

        loop {

            let state = harness.mirror.snapshot();

            if state.mentions.is_some() {

                break state.user_id.clone().unwrap();

            }

            updates.next().await.unwrap();

        }

    })

    .await

    .expect("the initial mentions snapshot did not arrive");

    let mentions_update = |account, revision| {

        Event::MentionsChanged(MentionsUpdate {

            snapshot: MentionsPage {

                entries: Vec::new(),

                next_cursor: None,

                end_reached: true,

                coverage: MentionCoverage::default(),

                unread_count: 0,

                revision,

                boundary: MentionsBoundary { account, revision },

            },

            upserts: Vec::new(),

            removed: Vec::new(),

        })

    };

    harness

        .inject

        .send(mentions_update(account.clone(), 100))

        .unwrap();

    next_matching(

        &mut updates,

        |change| matches!(change, StateChange::Mentions(update) if update.snapshot.revision == 100),

    )

    .await;

    assert_eq!(

        harness

            .mirror

            .snapshot()

            .mentions

            .as_ref()

            .unwrap()

            .revision,

        100

    );

    harness