laburnum 1.17.0

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

use {
  super::storage::{
    TestPartition,
    TestPartitions,
    TestRecordData,
  },
  crate::{
    Ident,
    database::{
      Database,
      chunk::RecordWriter,
      PartitionKey,
      query::QueryClient,
    },
    record::LaburnumRecord,
  },
  macro_rules_attribute::apply,
  smol_macros::test,
  std::collections::HashSet,
};

fn make_test_record(value: &str) -> TestRecordData {
  TestRecordData::Laburnum(LaburnumRecord::WorkspaceConfig {
    value: value.to_string(),
  })
}

#[apply(test!)]
async fn test_query_single_record_no_duplicates() {
  let db = Database::<TestPartitions>::new();
  let source_cache = crate::source::cache::reporter::SourceCacheReader::new_empty_for_test();

  let mut writer = RecordWriter::new(Ident::new("task1"));
  writer.insert::<TestPartition, _>("sk1", make_test_record("value1"));
  db.commit_chunk(writer.build(), &source_cache);

  let mut client = QueryClient::new(db);
  let results = client
    .prefix_internal(TestPartition::KEY, String::new())
    .await;

  assert_eq!(results.len(), 1, "Should return exactly 1 record");
}

#[apply(test!)]
async fn test_query_many_chunks_same_partition_different_sort_keys() {
  let db = Database::<TestPartitions>::new();
  let source_cache = crate::source::cache::reporter::SourceCacheReader::new_empty_for_test();

  for i in 0..15 {
    let task_id = Ident::new(&format!("task_{}", i));
    let mut writer = RecordWriter::new(task_id);
    writer.insert::<TestPartition, _>(
      format!("sort_key_{}", i),
      make_test_record(&format!("value_{}", i)),
    );
    db.commit_chunk(writer.build(), &source_cache);
  }

  let mut client = QueryClient::new(db.clone());
  let results = client
    .prefix_internal(TestPartition::KEY, String::new())
    .await;

  assert_eq!(results.len(), 15, "Should return exactly 15 records");

  let sort_keys: HashSet<_> =
    results.records().iter().map(|r| &r.sort_key).collect();
  assert_eq!(sort_keys.len(), 15, "All sort keys should be unique");
}

#[apply(test!)]
async fn test_query_multiple_chunks_same_sort_key() {
  let db = Database::<TestPartitions>::new();
  let source_cache = crate::source::cache::reporter::SourceCacheReader::new_empty_for_test();
  let sk = "same_sort_key";

  for i in 0..10 {
    let task_id = Ident::new(&format!("task_{}", i));
    let mut writer = RecordWriter::new(task_id);
    writer.insert::<TestPartition, _>(sk, make_test_record(&format!("value_{}", i)));
    db.commit_chunk(writer.build(), &source_cache);
  }

  let mut client = QueryClient::new(db.clone());
  let results = client
    .prefix_internal(TestPartition::KEY, String::new())
    .await;

  eprintln!(
    "test_query_multiple_chunks_same_sort_key: results.len() = {}",
    results.len()
  );

  for record in results.records() {
    eprintln!(
      "  pk={:?}, sk={}",
      record.partition_key, record.sort_key,
    );
  }
}

fn setup_range_test_db() -> Database<TestPartitions> {
  let db = Database::<TestPartitions>::new();
  let source_cache = crate::source::cache::reporter::SourceCacheReader::new_empty_for_test();

  let mut writer = RecordWriter::new(Ident::new("task1"));
  for key in ["a", "b", "c", "d"] {
    writer.insert::<TestPartition, _>(key, make_test_record(key));
  }
  db.commit_chunk(writer.build(), &source_cache);
  db
}

fn sort_keys_from_results(results: &crate::database::query_results::QueryResults<TestPartitions>) -> Vec<String> {
  results.records().iter().map(|r| r.sort_key.clone()).collect()
}

#[apply(test!)]
async fn test_index_less_than_exclusive() {
  let db = setup_range_test_db();
  let mut client = QueryClient::new(db);
  let results = client.less_than_internal(TestPartition::KEY, "c".to_string(), false).await;
  assert_eq!(sort_keys_from_results(&results), vec!["a", "b"]);
}

#[apply(test!)]
async fn test_index_less_than_inclusive() {
  let db = setup_range_test_db();
  let mut client = QueryClient::new(db);
  let results = client.less_than_internal(TestPartition::KEY, "c".to_string(), true).await;
  assert_eq!(sort_keys_from_results(&results), vec!["a", "b", "c"]);
}

#[apply(test!)]
async fn test_index_greater_than_exclusive() {
  let db = setup_range_test_db();
  let mut client = QueryClient::new(db);
  let results = client.greater_than_internal(TestPartition::KEY, "b".to_string(), false).await;
  assert_eq!(sort_keys_from_results(&results), vec!["c", "d"]);
}

#[apply(test!)]
async fn test_index_greater_than_inclusive() {
  let db = setup_range_test_db();
  let mut client = QueryClient::new(db);
  let results = client.greater_than_internal(TestPartition::KEY, "b".to_string(), true).await;
  assert_eq!(sort_keys_from_results(&results), vec!["b", "c", "d"]);
}

#[apply(test!)]
async fn test_index_between() {
  let db = setup_range_test_db();
  let mut client = QueryClient::new(db);
  let results = client.between_internal(TestPartition::KEY, "b".to_string(), "c".to_string()).await;
  assert_eq!(sort_keys_from_results(&results), vec!["b", "c"]);
}

#[apply(test!)]
async fn test_index_less_than_empty() {
  let db = setup_range_test_db();
  let mut client = QueryClient::new(db);
  let results = client.less_than_internal(TestPartition::KEY, "a".to_string(), false).await;
  assert_eq!(results.len(), 0);
}

#[apply(test!)]
async fn test_index_greater_than_empty() {
  let db = setup_range_test_db();
  let mut client = QueryClient::new(db);
  let results = client.greater_than_internal(TestPartition::KEY, "d".to_string(), false).await;
  assert_eq!(results.len(), 0);
}