238 lines
8.5 KiB
C++
238 lines
8.5 KiB
C++
// Rich test data generation utilities shared across the test suite.
|
|
#pragma once
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <chrono>
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <iomanip>
|
|
#include <map>
|
|
#include <random>
|
|
#include <sstream>
|
|
#include <string>
|
|
#include <string_view>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include <yams/metadata/document_metadata.h>
|
|
|
|
namespace yams::test {
|
|
|
|
class TestDataGenerator {
|
|
public:
|
|
TestDataGenerator() : TestDataGenerator(std::random_device{}()) {}
|
|
|
|
explicit TestDataGenerator(uint64_t seed) : rng_(seed) {}
|
|
|
|
[[nodiscard]] std::string generateTextDocument(std::size_t bytes,
|
|
std::string_view topic = "yams") {
|
|
static constexpr std::array words{
|
|
"search", "daemon", "vector", "chunk", "metadata",
|
|
"storage", "content", "index", "retrieval", "pipeline",
|
|
"ingest", "stream", "compression", "wal", "fixture",
|
|
"builder", "document", "embedding", "semantic", "analysis",
|
|
"orchestrator", "coordinator", "backoff", "resilient", "replay",
|
|
};
|
|
|
|
if (bytes == 0) {
|
|
return {};
|
|
}
|
|
|
|
std::uniform_int_distribution<std::size_t> word_dist(0, words.size() - 1);
|
|
std::ostringstream oss;
|
|
oss << "# " << topic << " log\n";
|
|
while (oss.tellp() < static_cast<std::streamoff>(bytes)) {
|
|
oss << words[word_dist(rng_)] << ' ';
|
|
if (std::bernoulli_distribution(0.15)(rng_)) {
|
|
oss << topic << ' ';
|
|
}
|
|
}
|
|
|
|
std::string result = oss.str();
|
|
if (result.size() > bytes) {
|
|
result.resize(bytes);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
[[nodiscard]] std::string generateMarkdown(int sections,
|
|
std::string_view title = "Generated Document") {
|
|
if (sections <= 0) {
|
|
sections = 1;
|
|
}
|
|
|
|
std::ostringstream oss;
|
|
oss << "# " << title << "\n\n";
|
|
for (int i = 0; i < sections; ++i) {
|
|
oss << "## Section " << (i + 1) << "\n";
|
|
oss << generateTextDocument(256, "section" + std::to_string(i + 1)) << "\n\n";
|
|
oss << "- item A" << i << "\n";
|
|
oss << "- item B" << i << "\n\n";
|
|
}
|
|
return oss.str();
|
|
}
|
|
|
|
[[nodiscard]] std::vector<std::byte> generateRandomBytes(std::size_t count) {
|
|
std::vector<std::byte> data(count);
|
|
std::uniform_int_distribution<int> dist(0, 255);
|
|
std::generate(data.begin(), data.end(), [&] { return std::byte(dist(rng_)); });
|
|
return data;
|
|
}
|
|
|
|
[[nodiscard]] std::vector<std::byte> generatePDF(std::size_t pages) {
|
|
pages = std::max<std::size_t>(pages, 1);
|
|
|
|
const std::size_t font_index = 3 + pages * 2;
|
|
const std::size_t total_objects = 3 + pages * 2;
|
|
|
|
std::string pdf;
|
|
pdf.reserve(2048 * pages);
|
|
pdf.append("%PDF-1.4\n");
|
|
|
|
std::vector<std::size_t> offsets;
|
|
offsets.reserve(total_objects + 1);
|
|
offsets.push_back(0U);
|
|
|
|
auto add_object = [&](const std::string& object) {
|
|
offsets.push_back(pdf.size());
|
|
pdf.append(object);
|
|
};
|
|
|
|
add_object("1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n");
|
|
|
|
std::ostringstream pages_stream;
|
|
pages_stream << "2 0 obj\n<< /Type /Pages /Kids [";
|
|
for (std::size_t i = 0; i < pages; ++i) {
|
|
if (i > 0) {
|
|
pages_stream << ' ';
|
|
}
|
|
pages_stream << (3 + i) << " 0 R";
|
|
}
|
|
pages_stream << "] /Count " << pages << " >>\nendobj\n";
|
|
add_object(pages_stream.str());
|
|
|
|
for (std::size_t i = 0; i < pages; ++i) {
|
|
const std::size_t page_object_index = 3 + i;
|
|
const std::size_t content_index = 3 + pages + i;
|
|
|
|
std::ostringstream page_obj;
|
|
page_obj << page_object_index
|
|
<< " 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] "
|
|
<< "/Resources << /Font << /F1 " << font_index << " 0 R >> >> /Contents "
|
|
<< content_index << " 0 R >>\nendobj\n";
|
|
add_object(page_obj.str());
|
|
|
|
std::string text = "Page " + std::to_string(i + 1) + " of " + std::to_string(pages);
|
|
std::ostringstream stream_obj;
|
|
stream_obj << "BT /F1 18 Tf 72 720 Td (" << text << ") Tj ET";
|
|
const std::string stream_payload = stream_obj.str();
|
|
|
|
std::ostringstream content_obj;
|
|
content_obj << content_index << " 0 obj\n<< /Length " << stream_payload.size()
|
|
<< " >>\nstream\n"
|
|
<< stream_payload << "\nendstream\nendobj\n";
|
|
add_object(content_obj.str());
|
|
}
|
|
|
|
std::ostringstream font_obj;
|
|
font_obj << font_index
|
|
<< " 0 obj\n<< /Type /Font /Subtype /Type1 /Name /F1 /BaseFont /Helvetica >>\n"
|
|
<< "endobj\n";
|
|
add_object(font_obj.str());
|
|
|
|
const std::size_t xref_offset = pdf.size();
|
|
std::ostringstream xref;
|
|
xref << "xref\n0 " << (offsets.size()) << "\n";
|
|
xref << "0000000000 65535 f \n";
|
|
for (std::size_t i = 1; i < offsets.size(); ++i) {
|
|
xref << std::setw(10) << std::setfill('0') << offsets[i] << " 00000 n \n";
|
|
}
|
|
xref << "trailer\n<< /Size " << offsets.size() << " /Root 1 0 R >>\n";
|
|
xref << "startxref\n" << xref_offset << "\n%%EOF\n";
|
|
pdf.append(xref.str());
|
|
|
|
std::vector<std::byte> bytes(pdf.size());
|
|
std::transform(pdf.begin(), pdf.end(), bytes.begin(),
|
|
[](unsigned char c) { return std::byte(c); });
|
|
return bytes;
|
|
}
|
|
|
|
[[nodiscard]] std::vector<std::byte> generateCorruptedPDF() {
|
|
auto data = generatePDF(1);
|
|
if (data.size() > 16) {
|
|
data.resize(data.size() / 2); // truncate to break structure
|
|
}
|
|
if (!data.empty()) {
|
|
data[0] = std::byte{0x00};
|
|
}
|
|
return data;
|
|
}
|
|
|
|
[[nodiscard]] std::vector<std::pair<std::string, metadata::MetadataValue>>
|
|
generateMetadata(std::size_t count, std::string_view prefix = "field") {
|
|
std::vector<std::pair<std::string, metadata::MetadataValue>> fields;
|
|
fields.reserve(count);
|
|
|
|
std::uniform_int_distribution<int> type_dist(0, 4);
|
|
std::uniform_int_distribution<int> int_dist(0, 10'000);
|
|
std::uniform_real_distribution<double> real_dist(0.0, 1.0);
|
|
std::bernoulli_distribution bool_dist(0.5);
|
|
|
|
for (std::size_t i = 0; i < count; ++i) {
|
|
const auto key = std::string(prefix) + "_" + std::to_string(i);
|
|
switch (type_dist(rng_)) {
|
|
case 0:
|
|
fields.emplace_back(key, metadata::MetadataValue("value_" + key));
|
|
break;
|
|
case 1:
|
|
fields.emplace_back(
|
|
key, metadata::MetadataValue(static_cast<int64_t>(int_dist(rng_))));
|
|
break;
|
|
case 2:
|
|
fields.emplace_back(key, metadata::MetadataValue(real_dist(rng_)));
|
|
break;
|
|
case 3:
|
|
fields.emplace_back(key, metadata::MetadataValue(bool_dist(rng_)));
|
|
break;
|
|
default:
|
|
fields.emplace_back(key, metadata::MetadataValue("meta:" + key));
|
|
break;
|
|
}
|
|
}
|
|
|
|
return fields;
|
|
}
|
|
|
|
[[nodiscard]] std::vector<std::string>
|
|
generateCorpus(std::size_t count, std::size_t averageSizeBytes, double duplicationRate = 0.1) {
|
|
std::vector<std::string> corpus;
|
|
corpus.reserve(count);
|
|
|
|
if (averageSizeBytes == 0) {
|
|
averageSizeBytes = 256;
|
|
}
|
|
|
|
std::poisson_distribution<int> size_dist(
|
|
static_cast<double>(std::max<std::size_t>(64, averageSizeBytes)));
|
|
std::bernoulli_distribution duplicate_dist(std::clamp(duplicationRate, 0.0, 1.0));
|
|
|
|
for (std::size_t i = 0; i < count; ++i) {
|
|
if (!corpus.empty() && duplicate_dist(rng_)) {
|
|
std::uniform_int_distribution<std::size_t> pick(0, corpus.size() - 1);
|
|
corpus.push_back(corpus[pick(rng_)]);
|
|
} else {
|
|
const auto target_size =
|
|
std::min<std::size_t>(size_dist(rng_), averageSizeBytes * 2);
|
|
corpus.push_back(generateTextDocument(target_size));
|
|
}
|
|
}
|
|
|
|
return corpus;
|
|
}
|
|
|
|
private:
|
|
std::mt19937_64 rng_;
|
|
};
|
|
|
|
} // namespace yams::test
|