laburnum 1.17.3

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

use {
  super::storage::{
    TestPartition,
    TestPartitions,
    TestRecordData,
    TestSortKey,
  },
  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".to_string(), make_test_record("value1"));
  db.commit_chunk(writer.build(), &source_cache);

  let mut client = QueryClient::new(db);
  let results = client
    .prefix_internal(TestPartition::KEY, TestSortKey::Test(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, TestSortKey::Test(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.to_string(), 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, TestSortKey::Test(String::new()))
    .await;

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

  for record in results.records() {
    let sk = match &record.sort_key {
      TestSortKey::Test(s) => s.as_str(),
      _ => "<unknown>",
    };
    eprintln!(
      "  pk={:?}, sk={}",
      record.partition_key, sk,
    );
  }
}

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.to_string(), 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| {
    match &r.sort_key {
      TestSortKey::Test(s) => s.clone(),
      _ => panic!("unexpected sort key variant"),
    }
  }).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, TestSortKey::Test("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, TestSortKey::Test("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, TestSortKey::Test("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, TestSortKey::Test("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, TestSortKey::Test("b".to_string()), TestSortKey::Test("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, TestSortKey::Test("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, TestSortKey::Test("d".to_string()), false).await;
  assert_eq!(results.len(), 0);
}