upstream-aw-server-rust/aw-datastore/tests/datastore.rs

#[macro_use]
extern crate log;
extern crate chrono;
#[macro_use]
extern crate aw_datastore;
extern crate serde_json;

extern crate dirs;

#[cfg(test)]
mod datastore_tests {
    use chrono::DateTime;
    use chrono::Duration;
    use chrono::Utc;
    use serde_json::json;

    use aw_datastore::Datastore;

    use aw_models::Bucket;
    use aw_models::BucketMetadata;
    use aw_models::Event;

    fn test_bucket() -> Bucket {
        Bucket {
            bid: None,
            id: "testid".to_string(),
            _type: "testtype".to_string(),
            client: "testclient".to_string(),
            hostname: "testhost".to_string(),
            created: None,
            data: json_map! {},
            metadata: BucketMetadata::default(),
            events: None,
            last_updated: None,
        }
    }

    fn create_test_bucket(ds: &Datastore) -> Bucket {
        let bucket = test_bucket();
        ds.create_bucket(&bucket).unwrap();
        bucket
    }

    fn create_named_test_bucket(ds: &Datastore, id: &str) -> Bucket {
        let mut bucket = test_bucket();
        bucket.id = id.to_string();
        ds.create_bucket(&bucket).unwrap();
        bucket
    }

    fn test_event(timestamp: chrono::DateTime<Utc>, duration: Duration) -> Event {
        Event {
            id: None,
            timestamp,
            duration,
            data: json_map! {"key": json!("value")},
        }
    }

    #[test]
    fn interval_access_paths_preserve_overlaps_ties_and_clipping() {
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);
        let epoch = chrono::DateTime::from_timestamp(1_000_000, 0).unwrap();
        let events: Vec<_> = [
            (0, 1),
            (1, 999),
            (900, 5),
            (900, 1),
            (900, 1),
            (901, 0),
            (1000, 1),
        ]
        .into_iter()
        .map(|(start, duration)| {
            test_event(
                epoch + Duration::seconds(start),
                Duration::seconds(duration),
            )
        })
        .collect();
        ds.insert_events(&bucket.id, &events).unwrap();
        for (start, end) in [(0, 10), (500, 510), (900, 905), (1100, 1200)] {
            let start = Some(epoch + Duration::seconds(start));
            let end = Some(epoch + Duration::seconds(end));
            // Limited reads always use the original starttime-first index;
            // unrestricted reads may use the endtime-first path.
            let original = ds.get_events(&bucket.id, start, end, Some(100)).unwrap();
            let selective = ds.get_events(&bucket.id, start, end, None).unwrap();
            assert_eq!(
                serde_json::to_value(&original).unwrap(),
                serde_json::to_value(&selective).unwrap()
            );
            assert_eq!(
                ds.get_event_count(&bucket.id, start, end).unwrap(),
                original.len() as i64
            );
            for limit in [0, 1, 3] {
                let limited = ds.get_events(&bucket.id, start, end, Some(limit)).unwrap();
                assert_eq!(
                    serde_json::to_value(limited).unwrap(),
                    serde_json::to_value(&original[..original.len().min(limit as usize)]).unwrap()
                );
            }
        }
    }

    #[test]
    fn test_migrate_test_bucket_names_renames_bucket_and_preserves_events() {
        let ds = Datastore::new_in_memory(false);
        let old_id = "aw-watcher-android-test_phone";
        let new_id = "aw-watcher-android_phone";
        create_named_test_bucket(&ds, old_id);
        let event = test_event(Utc::now(), Duration::seconds(30));
        ds.insert_events(old_id, std::slice::from_ref(&event))
            .unwrap();

        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 1);
        assert!(!ds.get_buckets().unwrap().contains_key(old_id));
        assert!(ds.get_buckets().unwrap().contains_key(new_id));
        assert_eq!(ds.get_events(new_id, None, None, None).unwrap().len(), 1);
    }

    #[test]
    fn test_migrate_test_bucket_names_merges_non_overlapping_buckets() {
        let ds = Datastore::new_in_memory(false);
        let old_id = "aw-watcher-android-test_phone";
        let new_id = "aw-watcher-android_phone";
        create_named_test_bucket(&ds, old_id);
        create_named_test_bucket(&ds, new_id);
        let now = Utc::now();
        ds.insert_events(
            old_id,
            &[test_event(now - Duration::hours(2), Duration::minutes(30))],
        )
        .unwrap();
        ds.insert_events(new_id, &[test_event(now, Duration::minutes(30))])
            .unwrap();

        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 1);
        assert!(!ds.get_buckets().unwrap().contains_key(old_id));
        assert_eq!(ds.get_events(new_id, None, None, None).unwrap().len(), 2);
    }

    #[test]
    fn test_migrate_test_bucket_names_keeps_overlapping_buckets_separate() {
        let ds = Datastore::new_in_memory(false);
        let old_id = "aw-watcher-android-test_phone";
        let new_id = "aw-watcher-android_phone";
        create_named_test_bucket(&ds, old_id);
        create_named_test_bucket(&ds, new_id);
        let now = Utc::now();
        ds.insert_events(old_id, &[test_event(now, Duration::minutes(30))])
            .unwrap();
        ds.insert_events(
            new_id,
            &[test_event(
                now + Duration::minutes(15),
                Duration::minutes(30),
            )],
        )
        .unwrap();

        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 0);
        assert!(ds.get_buckets().unwrap().contains_key(old_id));
        assert_eq!(ds.get_events(old_id, None, None, None).unwrap().len(), 1);
        assert_eq!(ds.get_events(new_id, None, None, None).unwrap().len(), 1);
    }

    #[test]
    fn test_migrate_test_bucket_names_moves_disjoint_events_when_some_overlap() {
        let ds = Datastore::new_in_memory(false);
        let old_id = "aw-watcher-android-test_phone";
        let new_id = "aw-watcher-android_phone";
        create_named_test_bucket(&ds, old_id);
        create_named_test_bucket(&ds, new_id);
        let now = Utc::now();
        ds.insert_events(
            old_id,
            &[
                test_event(now - Duration::hours(2), Duration::minutes(30)),
                test_event(now, Duration::minutes(30)),
            ],
        )
        .unwrap();
        ds.insert_events(
            new_id,
            &[test_event(
                now + Duration::minutes(15),
                Duration::minutes(30),
            )],
        )
        .unwrap();

        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 0);
        assert!(ds.get_buckets().unwrap().contains_key(old_id));
        assert_eq!(ds.get_events(old_id, None, None, None).unwrap().len(), 1);
        assert_eq!(ds.get_events(new_id, None, None, None).unwrap().len(), 2);
    }

    #[test]
    fn test_migrate_test_bucket_names_keeps_overlapping_legacy_events_together() {
        let ds = Datastore::new_in_memory(false);
        let old_id = "aw-watcher-android-test_phone";
        let new_id = "aw-watcher-android_phone";
        create_named_test_bucket(&ds, old_id);
        create_named_test_bucket(&ds, new_id);
        let now = Utc::now();
        ds.insert_events(
            old_id,
            &[
                test_event(now - Duration::hours(3), Duration::minutes(20)),
                test_event(now, Duration::minutes(30)),
                test_event(now + Duration::minutes(15), Duration::minutes(30)),
            ],
        )
        .unwrap();
        ds.insert_events(
            new_id,
            &[test_event(now + Duration::hours(2), Duration::minutes(20))],
        )
        .unwrap();

        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 0);
        assert!(ds.get_buckets().unwrap().contains_key(old_id));
        assert_eq!(ds.get_events(old_id, None, None, None).unwrap().len(), 2);
        assert_eq!(ds.get_events(new_id, None, None, None).unwrap().len(), 2);
    }

    #[test]
    fn test_migrate_test_bucket_names_zero_duration_event_at_shared_start_stays() {
        // Strict overlap semantics: a zero-duration legacy event that starts exactly
        // where a destination event starts is inside that event's open interval.
        let ds = Datastore::new_in_memory(false);
        let old_id = "aw-watcher-android-test_phone";
        let new_id = "aw-watcher-android_phone";
        create_named_test_bucket(&ds, old_id);
        create_named_test_bucket(&ds, new_id);
        let now = Utc::now();
        ds.insert_events(
            old_id,
            &[test_event(now + Duration::minutes(1), Duration::zero())],
        )
        .unwrap();
        ds.insert_events(new_id, &[test_event(now, Duration::minutes(30))])
            .unwrap();

        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 0);
        assert_eq!(ds.get_events(old_id, None, None, None).unwrap().len(), 1);
        assert_eq!(ds.get_events(new_id, None, None, None).unwrap().len(), 1);
    }

    #[test]
    fn test_migrate_test_bucket_names_merges_large_history_quickly() {
        // Regression for ActivityWatch/aw-android#261: the merge used a correlated
        // overlap subquery per legacy event (O(n^2)), which took hours on a phone
        // with a couple of years of history and wedged the datastore worker. The
        // sorted sweep must handle six-figure histories in seconds.
        let ds = Datastore::new_in_memory(false);
        let old_id = "aw-watcher-android-test_phone";
        let new_id = "aw-watcher-android_phone";
        create_named_test_bucket(&ds, old_id);
        create_named_test_bucket(&ds, new_id);
        let legacy_count = 100_000;
        let destination_count = 10_000;
        let start = Utc::now() - Duration::days(400);
        let step = Duration::minutes(5);
        let legacy: Vec<Event> = (0..legacy_count)
            .map(|i| test_event(start + step * i, step - Duration::seconds(1)))
            .collect();
        for chunk in legacy.chunks(10_000) {
            ds.insert_events(old_id, chunk).unwrap();
        }
        let cutover = start + step * legacy_count;
        // One overlapping cutover event must stay behind while the rest merges.
        ds.insert_events(
            old_id,
            &[test_event(
                cutover - Duration::minutes(2),
                Duration::minutes(4),
            )],
        )
        .unwrap();
        let destination: Vec<Event> = (0..destination_count)
            .map(|i| test_event(cutover + step * i, step - Duration::seconds(1)))
            .collect();
        ds.insert_events(new_id, &destination).unwrap();

        let started = std::time::Instant::now();
        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 0);
        let elapsed = started.elapsed();
        assert!(
            elapsed < std::time::Duration::from_secs(30),
            "merge took {elapsed:?}; the overlap scan regressed to quadratic time"
        );
        assert_eq!(
            ds.get_event_count(old_id, None, None).unwrap(),
            2,
            "only the overlapping cutover pair stays in the legacy bucket"
        );
        // Everything except the legacy event that the cutover overlaps moved over.
        assert_eq!(
            ds.get_event_count(new_id, None, None).unwrap(),
            (legacy_count - 1 + destination_count) as i64,
        );
    }

    #[test]
    fn test_migrate_test_bucket_names_heavy_overlap_is_linear() {
        // Stacked histories must not reintroduce a quadratic open-set scan.
        // Each event starts 1s later and outlives every later start, so the
        // open heap grows to n; scanning it per event would be ~n².
        let ds = Datastore::new_in_memory(false);
        let old_id = "aw-watcher-android-test_phone";
        let new_id = "aw-watcher-android_phone";
        create_named_test_bucket(&ds, old_id);
        create_named_test_bucket(&ds, new_id);
        let n = 20_000;
        let start = Utc::now();
        let span = Duration::seconds(n as i64);
        let legacy: Vec<Event> = (0..n)
            .map(|i| test_event(start + Duration::seconds(i as i64), span))
            .collect();
        for chunk in legacy.chunks(5_000) {
            ds.insert_events(old_id, chunk).unwrap();
        }

        let started = std::time::Instant::now();
        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 0);
        let elapsed = started.elapsed();
        assert!(
            elapsed < std::time::Duration::from_secs(10),
            "heavy-overlap merge took {elapsed:?}; open-set scan is quadratic"
        );
        assert_eq!(
            ds.get_event_count(old_id, None, None).unwrap(),
            n as i64,
            "every stacked event overlaps another, so none should move"
        );
        assert_eq!(ds.get_event_count(new_id, None, None).unwrap(), 0);
    }

    #[test]
    fn test_migrate_test_bucket_names_merges_interleaved_non_overlapping_events() {
        let ds = Datastore::new_in_memory(false);
        let old_id = "aw-watcher-android-test_phone";
        let new_id = "aw-watcher-android_phone";
        create_named_test_bucket(&ds, old_id);
        create_named_test_bucket(&ds, new_id);
        let now = Utc::now();
        ds.insert_events(
            old_id,
            &[
                test_event(now - Duration::hours(2), Duration::minutes(20)),
                test_event(now, Duration::minutes(20)),
            ],
        )
        .unwrap();
        ds.insert_events(
            new_id,
            &[test_event(now - Duration::hours(1), Duration::minutes(20))],
        )
        .unwrap();

        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 1);
        assert!(!ds.get_buckets().unwrap().contains_key(old_id));
        assert_eq!(ds.get_events(new_id, None, None, None).unwrap().len(), 3);
    }

    #[test]
    fn test_key_values_empty_result() {
        let ds = Datastore::new_in_memory(false);

        let values = ds.get_key_values("settings.%").unwrap();

        assert!(values.is_empty());
    }

    #[test]
    fn test_bucket_create_delete() {
        // Setup datastore
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        // Fetch bucket
        let bucket_fetched = ds.get_bucket(&bucket.id).unwrap();
        assert_eq!(bucket_fetched.id, bucket.id);
        assert_eq!(bucket_fetched._type, bucket._type);
        assert_eq!(bucket_fetched.client, bucket.client);
        assert_eq!(bucket_fetched.hostname, bucket.hostname);
        assert_eq!(bucket_fetched.metadata.end, None);

        match bucket_fetched.created {
            None => panic!("Expected 'None' in bucket to be replaced with current time"),
            Some(created) => {
                let now = Utc::now();
                assert!(created <= now);
                assert!(created > now - Duration::seconds(10));
            }
        };

        // Fetch all buckets
        let fetched_buckets = ds.get_buckets().unwrap();
        assert!(fetched_buckets.contains_key(&bucket.id));
        assert_eq!(fetched_buckets[&bucket.id].id, bucket.id);
        assert_eq!(fetched_buckets[&bucket.id]._type, bucket._type);
        assert_eq!(fetched_buckets[&bucket.id].client, bucket.client);
        assert_eq!(fetched_buckets[&bucket.id].hostname, bucket.hostname);
        assert_eq!(
            fetched_buckets[&bucket.id].metadata.start,
            bucket.metadata.start
        );
        assert_eq!(
            fetched_buckets[&bucket.id].metadata.end,
            bucket.metadata.end
        );

        match fetched_buckets[&bucket.id].created {
            None => panic!("Expected 'None' in bucket to be replaced with current time"),
            Some(created) => {
                let now = Utc::now();
                assert!(created <= now);
                assert!(created > now - Duration::seconds(10));
            }
        };

        // Delete bucket
        match ds.delete_bucket(&bucket.id) {
            Ok(_) => info!("bucket successfully deleted"),
            Err(e) => panic!("{e:?}"),
        }
        match ds.get_bucket(&bucket.id) {
            Ok(_) => {
                panic!("Expected datastore to delete bucket but bucket seems to still be available")
            }
            Err(_e) => (),
        }
    }

    #[test]
    fn test_events_get_single() {
        // Setup datastore
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        // Insert event
        let e1 = Event {
            id: None,
            timestamp: Utc::now(),
            duration: Duration::seconds(0),
            data: json_map! {"key": json!("value")},
        };
        let mut e2 = e1.clone();
        e2.timestamp += Duration::nanoseconds(1);

        let event_list = [e1, e2];
        ds.insert_events(&bucket.id, &event_list).unwrap();

        let events = ds.get_events(&bucket.id, None, None, None).unwrap();
        let first_event = events.first().unwrap();
        let first_event_id = first_event.id.unwrap();

        let fetched_event = ds.get_event(&bucket.id, first_event_id).unwrap();
        // TODO: Check entire events to ensure integrity
        assert_eq!(fetched_event.id.unwrap(), first_event_id);
    }

    #[test]
    fn test_events_get_filters() {
        // Setup datastore
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        // Insert event
        let e1 = Event {
            id: None,
            timestamp: Utc::now(),
            duration: Duration::seconds(0),
            data: json_map! {"key": json!("value")},
        };
        let mut e2 = e1.clone();
        e2.timestamp += Duration::nanoseconds(1);

        let event_list = [e1.clone(), e2.clone()];

        ds.insert_events(&bucket.id, &event_list).unwrap();

        // Get all events
        let fetched_events_all = ds.get_events(&bucket.id, None, None, None).unwrap();
        let expected_fetched_events = [&e2, &e1];
        assert_eq!(fetched_events_all.len(), 2);
        for i in 0..fetched_events_all.len() {
            let expected = &expected_fetched_events[i];
            let new = &fetched_events_all[i];
            assert_eq!(new.timestamp, expected.timestamp);
            assert_eq!(new.duration, expected.duration);
            assert_eq!(new.data, expected.data);
        }

        info!("Get events with limit filter");
        let fetched_events_limit = ds.get_events(&bucket.id, None, None, Some(1)).unwrap();
        assert_eq!(fetched_events_limit.len(), 1);
        assert_eq!(fetched_events_limit[0].timestamp, e2.timestamp);
        assert_eq!(fetched_events_limit[0].duration, e2.duration);
        assert_eq!(fetched_events_limit[0].data, e2.data);

        info!("Get events with starttime filter");
        let fetched_events_start = ds
            .get_events(&bucket.id, Some(e2.timestamp), None, None)
            .unwrap();
        assert_eq!(fetched_events_start.len(), 1);
        assert_eq!(fetched_events_start[0].timestamp, e2.timestamp);
        assert_eq!(fetched_events_start[0].duration, e2.duration);
        assert_eq!(fetched_events_start[0].data, e2.data);

        info!("Get events with endtime filter");
        let fetched_events_start = ds
            .get_events(&bucket.id, None, Some(e1.timestamp), None)
            .unwrap();
        assert_eq!(fetched_events_start.len(), 1);
        assert_eq!(fetched_events_start[0].timestamp, e1.timestamp);
        assert_eq!(fetched_events_start[0].duration, e1.duration);
        assert_eq!(fetched_events_start[0].data, e1.data);

        // Get eventcount
        let event_count = ds.get_event_count(&bucket.id, None, None).unwrap();
        assert_eq!(event_count, 2);
    }

    /// Tests that events that cover a timeperiod get included when that timeperiod is queried.
    #[test]
    fn test_get_events_filters_cover() {
        // TODO: Also test event-cutoff, although perhaps that happens in the transforms/queries?

        // Setup datastore
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        let now = Utc::now();

        // Insert event
        let e1 = Event {
            id: None,
            timestamp: Utc::now(),
            duration: Duration::seconds(100),
            data: json_map! {"key": json!("value")},
        };

        let event_list = [e1];
        ds.insert_events(&bucket.id, &event_list).unwrap();

        info!("Get event that covers queried timeperiod");
        let query_start = now + Duration::seconds(1);
        let query_end = query_start + Duration::seconds(1);
        let fetched_events_limit = ds
            .get_events(&bucket.id, Some(query_start), Some(query_end), Some(1))
            .unwrap();
        assert_eq!(fetched_events_limit.len(), 1);

        // Get eventcount
        let event_count = ds
            .get_event_count(&bucket.id, Some(query_start), Some(query_end))
            .unwrap();
        assert_eq!(event_count, 1);
    }

    #[test]
    fn test_events_delete() {
        // Setup datastore
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        // Insert event
        let e1 = Event {
            id: None,
            timestamp: Utc::now(),
            duration: Duration::seconds(0),
            data: json_map! {"key": json!("value")},
        };
        let mut e2 = e1.clone();
        e2.timestamp += Duration::seconds(1);

        let event_list = [e1.clone(), e2.clone()];

        ds.insert_events(&bucket.id, &event_list).unwrap();

        // Get all events
        let fetched_events_all = ds.get_events(&bucket.id, None, None, None).unwrap();
        let expected_fetched_events = [&e2, &e1];
        assert_eq!(fetched_events_all.len(), 2);
        for i in 0..fetched_events_all.len() {
            let expected = &expected_fetched_events[i];
            let new = &fetched_events_all[i];
            assert_eq!(new.timestamp, expected.timestamp);
            assert_eq!(new.duration, expected.duration);
            assert_eq!(new.data, expected.data);
        }
        let e1 = &fetched_events_all[0];
        let e2 = &fetched_events_all[1];

        // Delete one event
        ds.delete_events_by_id(&bucket.id, vec![e1.id.unwrap()])
            .unwrap();

        // Get all events
        let fetched_events_all = ds.get_events(&bucket.id, None, None, None).unwrap();
        let expected_fetched_events = [e2];
        assert_eq!(fetched_events_all.len(), 1);
        for i in 0..fetched_events_all.len() {
            let expected = &expected_fetched_events[i];
            let new = &fetched_events_all[i];
            assert_eq!(new.id, expected.id);
            assert_eq!(new.timestamp, expected.timestamp);
            assert_eq!(new.duration, expected.duration);
            assert_eq!(new.data, expected.data);
        }
    }

    #[test]
    fn test_events_delete_refreshes_bounds_with_pre_epoch_subsecond_event() {
        // Regression test: refreshing bucket bounds after a delete must not
        // panic when a remaining event has a pre-epoch timestamp with a
        // subsecond component (negative nanoseconds, not a multiple of 1s).
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        // 1969-12-31T23:59:55.250Z: negative seconds, with a fractional part.
        let pre_epoch = DateTime::from_timestamp(-5, 250_000_000).unwrap();
        let e_pre_epoch = Event {
            id: None,
            timestamp: pre_epoch,
            duration: Duration::seconds(0),
            data: json_map! {"key": json!("pre-epoch")},
        };
        let e_other = Event {
            id: None,
            timestamp: Utc::now(),
            duration: Duration::seconds(0),
            data: json_map! {"key": json!("other")},
        };

        // get_events with an unbounded (None) starttime defaults to the
        // epoch, which would silently exclude the pre-epoch event - so fetch
        // the inserted events' assigned ids directly instead.
        let event_list = [e_pre_epoch.clone(), e_other.clone()];
        let inserted = ds.insert_events(&bucket.id, &event_list).unwrap();
        let inserted_other = inserted
            .iter()
            .find(|e| e.timestamp == e_other.timestamp)
            .unwrap();

        // Deleting the other event forces refresh_bucket_bounds to recompute
        // start/end from the remaining pre-epoch event; this must not panic.
        ds.delete_events_by_id(&bucket.id, vec![inserted_other.id.unwrap()])
            .unwrap();

        let bucket_fetched = ds.get_bucket(&bucket.id).unwrap();
        assert_eq!(bucket_fetched.metadata.start, Some(pre_epoch));
        assert_eq!(
            bucket_fetched.metadata.end,
            Some(e_pre_epoch.calculate_endtime())
        );
    }

    #[test]
    fn test_bucket_metadata_start_end() {
        // Setup datastore
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        // Insert event
        let e1 = Event {
            id: None,
            timestamp: Utc::now(),
            duration: Duration::seconds(0),
            data: json_map! {"key": json!("value")},
        };
        let mut e2 = e1.clone();
        e2.timestamp += Duration::nanoseconds(1);

        let event_list = [e1.clone(), e2.clone()];

        ds.insert_events(&bucket.id, &event_list).unwrap();

        // Validate correct start and end in bucket
        let bucket_fetched = ds.get_bucket(&bucket.id).unwrap();
        assert_eq!(bucket_fetched.metadata.start, Some(e1.timestamp));
        assert_eq!(bucket_fetched.metadata.end, Some(e2.calculate_endtime()));
    }

    #[test]
    fn test_event_heartbeat() {
        // Setup datastore
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        // Insert event
        let e1 = Event {
            id: None,
            timestamp: Utc::now(),
            duration: Duration::seconds(0),
            data: json_map! {"key": json!("value")},
        };
        let mut e2 = e1.clone();
        e2.timestamp += Duration::seconds(1);

        let mut e_diff_data = e2.clone();
        e_diff_data.timestamp += Duration::seconds(1);
        e_diff_data.data = json_map! {"key": json!("other value")};

        // First event
        ds.heartbeat(&bucket.id, e1.clone(), 10.0).unwrap();
        let fetched_events = ds.get_events(&bucket.id, None, None, None).unwrap();
        assert_eq!(fetched_events.len(), 1);
        assert_eq!(fetched_events[0].timestamp, e1.timestamp);
        assert_eq!(fetched_events[0].duration, e1.duration);
        assert_eq!(fetched_events[0].data, e1.data);
        let e1 = &fetched_events[0];

        // Heartbeat match
        ds.heartbeat(&bucket.id, e2, 10.0).unwrap();
        let fetched_events = ds.get_events(&bucket.id, None, None, None).unwrap();
        assert_eq!(fetched_events.len(), 1);
        assert_eq!(fetched_events[0].timestamp, e1.timestamp);
        assert_eq!(fetched_events[0].duration, Duration::seconds(1));
        assert_eq!(fetched_events[0].data, e1.data);
        assert_eq!(fetched_events[0].id, e1.id);
        let e2 = &fetched_events[0];

        // Heartbeat diff
        ds.heartbeat(&bucket.id, e_diff_data.clone(), 10.0).unwrap();
        let fetched_events = ds.get_events(&bucket.id, None, None, None).unwrap();
        assert_eq!(fetched_events.len(), 2);
        assert_eq!(fetched_events[0].timestamp, e_diff_data.timestamp);
        assert_eq!(fetched_events[0].duration, e_diff_data.duration);
        assert_eq!(fetched_events[0].data, e_diff_data.data);
        assert_ne!(fetched_events[0].id, e2.id);
    }

    #[test]
    fn test_event_heartbeat_consecutive_merges() {
        // Regression test for https://github.com/ActivityWatch/aw-server-rust/issues/559
        // After insert_heartbeat merges two events, the merged event stored in cache
        // must retain the DB event ID so subsequent heartbeats can find it for
        // replace_last_event. Without the ID, every other heartbeat would fail.
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        let now = Utc::now();
        let pulsetime = 10.0;

        // Send 5 consecutive heartbeats with the same data, 1 second apart.
        // They should all merge into a single event with increasing duration.
        for i in 0..5 {
            let e = Event {
                id: None,
                timestamp: now + Duration::seconds(i),
                duration: Duration::seconds(0),
                data: json_map! {"key": json!("value")},
            };
            ds.heartbeat(&bucket.id, e, pulsetime).unwrap();
        }

        let events = ds.get_events(&bucket.id, None, None, None).unwrap();
        assert_eq!(
            events.len(),
            1,
            "all heartbeats should merge into one event"
        );
        assert_eq!(events[0].timestamp, now);
        assert_eq!(events[0].duration, Duration::seconds(4));
        assert!(events[0].id.is_some(), "event should have a DB id");
    }

    #[test]
    fn test_event_replace() {
        // Setup datastore
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        // Insert event
        let e = Event {
            id: None,
            timestamp: Utc::now(),
            duration: Duration::seconds(0),
            data: json_map! {"key": json!("value")},
        };
        let mut e1 = e.clone();
        e1.data = json_map! {"key": json!("value1")};
        let mut e2 = e.clone();
        e2.data = json_map! {"key": json!("value2")};
        let mut e3 = e;
        e3.data = json_map! {"key": json!("value3")};

        let events_init = &[e1, e2, e3];
        let events_ret = ds.insert_events(&bucket.id, events_init).unwrap();
        // Validate return from insert
        assert_eq!(events_ret[0].id, Some(1));
        assert_eq!(events_ret[1].id, Some(2));
        assert_eq!(events_ret[2].id, Some(3));
        assert_eq!(events_ret.len(), 3);
        assert_eq!(events_ret[0], events_init[0]);
        assert_eq!(events_ret[1], events_init[1]);
        assert_eq!(events_ret[2], events_init[2]);

        let events_init = events_ret;

        // Insert e2 with identical data and id (which means a replace)
        {
            let events_ret = ds
                .insert_events(&bucket.id, &[events_init[1].clone()])
                .unwrap();
            assert_eq!(events_ret.len(), 1);
            assert_eq!(events_ret[0], events_init[1]);
            let fetched_events = ds.get_events(&bucket.id, None, None, None).unwrap();
            assert_eq!(fetched_events, events_init);
        }

        // Make new event with same id but different data
        {
            let mut e2 = events_init[1].clone();
            e2.data = json_map! {"key": json!("value2_modified")};
            let events_ret = ds.insert_events(&bucket.id, &[e2.clone()]).unwrap();
            assert_eq!(events_ret.len(), 1);
            assert_eq!(events_ret[0], e2);
            let fetched_events = ds.get_events(&bucket.id, None, None, None).unwrap();
            assert_eq!(fetched_events.len(), 3);
            assert_eq!(fetched_events[1], e2);
            assert_eq!(fetched_events[0].id, Some(1));
            assert_eq!(fetched_events[1].id, Some(2));
            assert_eq!(fetched_events[2].id, Some(3));
        }
    }

    #[test]
    fn test_migration_v4_to_v5() {
        let test_dir = tempfile::tempdir().unwrap();
        let db_path = test_dir.path().join("datastore-unittest-migration-v4.db");
        let db_path_str = db_path.to_str().unwrap().to_string();

        if db_path.exists() {
            std::fs::remove_file(&db_path)
                .expect("Failed to remove datastore-unittest-migration-v4.db file");
        }

        // Construct a database with the v4 schema (single-column indexes)
        {
            let conn = rusqlite::Connection::open(&db_path).unwrap();
            conn.execute_batch(
                r#"
                CREATE TABLE buckets (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT UNIQUE NOT NULL,
                    type TEXT NOT NULL,
                    client TEXT NOT NULL,
                    hostname TEXT NOT NULL,
                    created TEXT NOT NULL,
                    data TEXT NOT NULL DEFAULT '{}'
                );
                CREATE INDEX bucket_id_index ON buckets(id);
                CREATE TABLE events (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    bucketrow INTEGER NOT NULL,
                    starttime INTEGER NOT NULL,
                    endtime INTEGER NOT NULL,
                    data TEXT NOT NULL,
                    FOREIGN KEY (bucketrow) REFERENCES buckets(id)
                );
                CREATE INDEX events_bucketrow_index ON events(bucketrow);
                CREATE INDEX events_starttime_index ON events(starttime);
                CREATE INDEX events_endtime_index ON events(endtime);
                CREATE TABLE key_value (
                    key TEXT PRIMARY KEY,
                    value TEXT,
                    last_modified NUMBER NOT NULL
                );
                INSERT INTO buckets (name, type, client, hostname, created, data)
                    VALUES ('testid', 'testtype', 'testclient', 'testhost',
                            '2024-01-01T00:00:00+00:00', '{}');
                INSERT INTO events (bucketrow, starttime, endtime, data)
                    VALUES (1, 1000000000, 2000000000, '{"key": "value"}');
                PRAGMA user_version = 4;
            "#,
            )
            .unwrap();
        }

        // Opening the datastore migrates to the newest version
        {
            let ds = Datastore::new(db_path_str, false);
            let events = ds.get_events("testid", None, None, None).unwrap();
            assert_eq!(events.len(), 1);
            assert_eq!(events[0].data, json_map! {"key": json!("value")});
            ds.close();
        }

        // Verify version bump and index replacement
        {
            let conn = rusqlite::Connection::open(&db_path).unwrap();
            let version: i32 = conn
                .pragma_query_value(None, "user_version", |row| row.get(0))
                .unwrap();
            assert_eq!(version, 6);
            let old_indexes: i64 = conn
                .query_row(
                    "SELECT count(*) FROM sqlite_master WHERE type = 'index' AND name IN
                     ('events_bucketrow_index', 'events_starttime_index', 'events_endtime_index')",
                    [],
                    |row| row.get(0),
                )
                .unwrap();
            assert_eq!(old_indexes, 0, "single-column indexes should be dropped");
            let new_index: i64 = conn
                .query_row(
                    "SELECT count(*) FROM sqlite_master WHERE type = 'index'
                     AND name = 'events_bucketrow_starttime_endtime_index'",
                    [],
                    |row| row.get(0),
                )
                .unwrap();
            assert_eq!(new_index, 1, "composite index should exist");
            let end_index: i64 = conn
                .query_row(
                    "SELECT count(*) FROM sqlite_master WHERE type = 'index'
                 AND name = 'events_bucketrow_endtime_starttime_index'",
                    [],
                    |row| row.get(0),
                )
                .unwrap();
            assert_eq!(end_index, 1);
        }

        std::fs::remove_file(&db_path)
            .expect("Failed to remove datastore-unittest-migration-v4.db file");
    }

    #[test]
    fn test_datastore_reload() {
        // Create tmp datastore path
        let test_dir = tempfile::tempdir().unwrap();
        let db_path = test_dir.path().join("datastore-unittest.db");
        let db_path_str = db_path.to_str().unwrap().to_string();

        if db_path.exists() {
            std::fs::remove_file(db_path.clone())
                .expect("Failed to remove datastore-unittest.db file");
        }

        let empty_bucket = test_bucket();
        let mut populated_bucket = empty_bucket.clone();
        populated_bucket.id = "testid2".to_string();
        let e1 = Event {
            id: None,
            timestamp: Utc::now(),
            duration: Duration::seconds(0),
            data: json_map! {"key": json!("value")},
        };
        {
            // Initialize database and create buckets
            let ds = Datastore::new(db_path_str.clone(), false);
            ds.create_bucket(&empty_bucket).unwrap();
            ds.create_bucket(&populated_bucket).unwrap();
            // Insert event
            ds.insert_events(&populated_bucket.id, std::slice::from_ref(&e1))
                .unwrap();

            // Check that all cached bucket data is correct
            let buckets = ds.get_buckets().unwrap();
            assert_eq!(buckets[&empty_bucket.id].metadata.start, None);
            assert_eq!(buckets[&empty_bucket.id].metadata.end, None);
            assert_eq!(
                buckets[&populated_bucket.id].metadata.start,
                Some(e1.calculate_endtime())
            );
            assert_eq!(
                buckets[&populated_bucket.id].metadata.end,
                Some(e1.timestamp)
            );
            ds.force_commit().unwrap();
        }
        {
            // Load database again
            let ds = Datastore::new(db_path_str, false);
            // Check that all bucket data is correct after reload
            let buckets = ds.get_buckets().unwrap();
            assert_eq!(buckets[&empty_bucket.id].metadata.start, None);
            assert_eq!(buckets[&empty_bucket.id].metadata.end, None);
            assert_eq!(
                buckets[&populated_bucket.id].metadata.start,
                Some(e1.timestamp)
            );
            assert_eq!(
                buckets[&populated_bucket.id].metadata.end,
                Some(e1.calculate_endtime())
            );
        }
    }

    /// Test that an encrypted datastore can be created, written to, and reopened with the same key
    /// with data intact.
    #[test]
    #[cfg(any(feature = "encryption", feature = "encryption-vendored"))]
    fn test_encrypted_datastore_roundtrip() {
        use std::fs;
        let test_dir = tempfile::tempdir().unwrap();
        let dir = test_dir.path();
        let db_path = dir.join("test-encrypted.db").to_str().unwrap().to_string();
        // Clean up from previous runs
        let _ = fs::remove_file(&db_path);

        let key = "s3cr3t-p@ssw0rd".to_string();

        // Create and populate encrypted datastore
        {
            let ds = Datastore::new_encrypted(db_path.clone(), key.clone(), false);
            let bucket = create_test_bucket(&ds);
            let e = Event {
                id: None,
                timestamp: Utc::now(),
                duration: Duration::seconds(1),
                data: json_map! { "app": "test-encrypted" },
            };
            let inserted = ds.insert_events(&bucket.id, &[e]).unwrap();
            assert_eq!(inserted.len(), 1);
            ds.force_commit().unwrap();
            ds.close();
        }

        // Reopen with correct key — data must survive the roundtrip
        {
            let ds = Datastore::new_encrypted(db_path.clone(), key.clone(), false);
            let events = ds
                .get_events("testid", None, None, None)
                .expect("should read events from encrypted DB after reopen");
            assert_eq!(events.len(), 1, "expected 1 event after encrypted reopen");
            assert_eq!(events[0].data["app"], "test-encrypted");
            ds.close();
        }

        let _ = fs::remove_file(&db_path);
    }

    fn privacy_event(title: &str) -> Event {
        Event {
            id: None,
            timestamp: Utc::now(),
            duration: Duration::seconds(1),
            data: json_map! {"title": json!(title), "app": json!("Firefox")},
        }
    }

    /// Saving privacy_filters used to persist the JSON and never load it into
    /// the worker's engine (ActivityWatch/aw-server-rust#659). A drop rule
    /// must actually drop a matching insert without an explicit refresh call.
    #[test]
    fn test_privacy_filter_applies_after_setting_saved() {
        let ds = Datastore::new_in_memory(false);
        let bucket = create_test_bucket(&ds);

        let drop_secret =
            r#"[{"enabled":true,"field":"title","pattern":"(?i)secret","action":"drop"}]"#;
        ds.set_key_value("settings.privacy_filters", drop_secret)
            .unwrap();

        ds.insert_events(
            &bucket.id,
            &[privacy_event("my secret file"), privacy_event("readme.md")],
        )
        .unwrap();

        let events = ds.get_events(&bucket.id, None, None, None).unwrap();
        assert_eq!(
            events.len(),
            1,
            "drop rule should discard the matching event"
        );
        assert_eq!(events[0].data.get("title").unwrap(), "readme.md");

        ds.delete_key_value("settings.privacy_filters").unwrap();
        ds.insert_events(&bucket.id, &[privacy_event("another secret")])
            .unwrap();
        let events = ds.get_events(&bucket.id, None, None, None).unwrap();
        assert_eq!(
            events.len(),
            2,
            "deleting the setting should stop filtering"
        );
    }

    /// Rules must load at worker startup, not only after a later SetKeyValue.
    #[test]
    fn test_privacy_filter_survives_datastore_reload() {
        let test_dir = tempfile::tempdir().unwrap();
        let db_path = test_dir
            .path()
            .join("datastore-unittest-privacy-filters.db");
        let db_path_str = db_path.to_str().unwrap().to_string();

        let drop_secret =
            r#"[{"enabled":true,"field":"title","pattern":"(?i)secret","action":"drop"}]"#;
        {
            let ds = Datastore::new(db_path_str.clone(), false);
            create_test_bucket(&ds);
            ds.set_key_value("settings.privacy_filters", drop_secret)
                .unwrap();
            ds.force_commit().unwrap();
            ds.close();
        }
        {
            let ds = Datastore::new(db_path_str, false);
            ds.insert_events("testid", &[privacy_event("top secret notes")])
                .unwrap();
            let events = ds.get_events("testid", None, None, None).unwrap();
            assert!(
                events.is_empty(),
                "persisted drop rule must apply after reopen, got {events:?}"
            );
            ds.close();
        }

        // Windows can still hold the SQLite handle after close() while the
        // worker thread unwinds (ERROR_SHARING_VIOLATION). Cleanup is best-effort.
        let _ = std::fs::remove_file(&db_path);
        let _ = std::fs::remove_file(db_path.with_extension("db-wal"));
        let _ = std::fs::remove_file(db_path.with_extension("db-shm"));
    }
}