Produce test data sets

"""
Build a test data set from a schema and a seed.

Rung N0. Standard library only, no file written: the function returns a list of
rows that a test can hand straight to the code under test.

Determinism is the whole point. The same seed and the same schema give exactly
the same rows, on every machine, in both languages, for ever. That is what
makes a failing test replayable and a regression reproducible: the seed printed
next to a failure is enough to rebuild the data that caused it.

Which is why the generator is written out here instead of being taken from the
platform. The standard generators of Python and JavaScript produce different
sequences from the same seed, so a data set built with them cannot be handed
from one language to the other, nor compared between a back end and a front end.

Each cell is drawn from a hash of (seed, field, row) rather than from a running
stream. Adding a field to the schema therefore leaves every other column
untouched, instead of shifting the whole set by one draw.
"""

from datetime import date, timedelta

# FNV-1a, the same constants as the rest of the catalogue, so that a given
# string hashes identically wherever it is hashed.
FNV_OFFSET = 2166136261
FNV_PRIME = 16777619
MASK32 = 0xFFFFFFFF

UNIT = "\x1f"  # separates the parts of a cell key, and appears in none of them


def stable_hash(text: str) -> int:
    """
    FNV-1a on 32 bits.

    Not the built-in hash: that one is salted per process, so the same seed
    would give different data after every restart.
    """
    digest = FNV_OFFSET
    for char in text:
        digest = ((digest ^ ord(char)) * FNV_PRIME) & MASK32
    return digest


def draw(seed: str, field: str, row: int) -> int:
    """The single source of randomness: one 32-bit integer per cell."""
    return stable_hash(f"{seed}{UNIT}{field}{UNIT}{row}")


def _value(spec: dict, number: int, row: int):
    """Turn one drawn integer into one value that satisfies the field spec."""
    kind = spec["type"]
    if kind == "int":
        low, high = spec["min"], spec["max"]
        if low > high:
            raise ValueError(f"impossible range: min {low} is above max {high}")
        return low + number % (high - low + 1)
    if kind == "choice":
        values = spec["values"]
        if not values:
            raise ValueError("a choice field needs at least one value")
        return values[number % len(values)]
    if kind == "bool":
        # Percentages, not probabilities: an observed rate is read off a
        # dashboard as a percentage, and copied here as one.
        return number % 100 < spec.get("true_percent", 50)
    if kind == "date":
        span = spec.get("days", 1)
        return (date.fromisoformat(spec["start"]) + timedelta(days=number % span)).isoformat()
    if kind == "sequence":
        # Unique by construction, because an identifier that repeats turns a
        # test about duplicates into a test about the generator.
        rank = f"{row + spec.get('start', 1):0{spec.get('width', 4)}d}"
        return f"{spec.get('prefix', '')}{rank}{spec.get('suffix', '')}"
    raise ValueError(f"unknown field type {kind!r}")


def generate_rows(schema: dict, count: int, seed: str) -> list[dict]:
    """
    Return `count` rows, each field drawn independently from its own spec.

    `schema` maps a field name to a spec: {"type": "int", "min": …, "max": …},
    {"type": "choice", "values": […]}, {"type": "bool", "true_percent": …},
    {"type": "date", "start": "YYYY-MM-DD", "days": …} or
    {"type": "sequence", "prefix": …, "width": …, "suffix": …}.
    """
    return [
        {field: _value(spec, draw(seed, field, row), row) for field, spec in schema.items()}
        for row in range(count)
    ]

/**
 * Build a test data set from a schema and a seed.
 *
 * Rung N0. No dependency, no file written: the function returns an array of
 * rows that a test can hand straight to the code under test.
 *
 * Determinism is the whole point. The same seed and the same schema give
 * exactly the same rows, on every machine, in both languages, for ever. That
 * is what makes a failing test replayable and a regression reproducible: the
 * seed printed next to a failure is enough to rebuild the data that caused it.
 *
 * Which is why the generator is written out here instead of being taken from
 * the platform. `Math.random` cannot be seeded at all, and the standard
 * generators of Python and JavaScript produce different sequences from the
 * same seed, so a data set built with them cannot be handed from one language
 * to the other, nor compared between a back end and a front end.
 *
 * Each cell is drawn from a hash of (seed, field, row) rather than from a
 * running stream. Adding a field to the schema therefore leaves every other
 * column untouched, instead of shifting the whole set by one draw.
 */

// FNV-1a, the same constants as the rest of the catalogue, so that a given
// string hashes identically wherever it is hashed.
const FNV_OFFSET = 2166136261;
const FNV_PRIME = 16777619;

// Unit separator: it joins the parts of a cell key, and appears in none of them.
const UNIT = '\u001f';

const DAY = 24 * 60 * 60 * 1000;

/**
 * FNV-1a on 32 bits.
 *
 * `Math.imul` is what keeps this identical to the Python version: a plain `*`
 * on numbers this large loses precision past 2^53 and silently drifts.
 */
export function stableHash(text) {
  let digest = FNV_OFFSET;
  for (const char of text) {
    digest = Math.imul(digest ^ char.codePointAt(0), FNV_PRIME) >>> 0;
  }
  return digest;
}

/** The single source of randomness: one 32-bit integer per cell. */
export function draw(seed, field, row) {
  return stableHash(`${seed}${UNIT}${field}${UNIT}${row}`);
}

/** Turn one drawn integer into one value that satisfies the field spec. */
function value(spec, number, row) {
  switch (spec.type) {
    case 'int': {
      const { min, max } = spec;
      if (min > max) throw new RangeError(`impossible range: min ${min} is above max ${max}`);
      return min + (number % (max - min + 1));
    }
    case 'choice': {
      if (!spec.values?.length) throw new RangeError('a choice field needs at least one value');
      return spec.values[number % spec.values.length];
    }
    case 'bool':
      // Percentages, not probabilities: an observed rate is read off a
      // dashboard as a percentage, and copied here as one.
      return number % 100 < (spec.true_percent ?? 50);
    case 'date': {
      const start = Date.parse(`${spec.start}T00:00:00Z`);
      return new Date(start + (number % (spec.days ?? 1)) * DAY).toISOString().slice(0, 10);
    }
    case 'sequence': {
      // Unique by construction, because an identifier that repeats turns a
      // test about duplicates into a test about the generator.
      const rank = String(row + (spec.start ?? 1)).padStart(spec.width ?? 4, '0');
      return `${spec.prefix ?? ''}${rank}${spec.suffix ?? ''}`;
    }
    default:
      throw new RangeError(`unknown field type ${JSON.stringify(spec.type)}`);
  }
}

/**
 * Return `count` rows, each field drawn independently from its own spec.
 *
 * `schema` maps a field name to a spec: {type: 'int', min, max},
 * {type: 'choice', values}, {type: 'bool', true_percent},
 * {type: 'date', start: 'YYYY-MM-DD', days} or
 * {type: 'sequence', prefix, width, suffix}.
 */
export function generateRows(schema, count, seed) {
  const rows = [];
  for (let row = 0; row < count; row += 1) {
    const built = {};
    for (const [field, spec] of Object.entries(schema)) {
      built[field] = value(spec, draw(seed, field, row), row);
    }
    rows.push(built);
  }
  return rows;
}

"""
Sample a test data set from the distributions observed in production.

Rung N1. Standard library only, no file written, and still no network: the
distributions come from the caller, as the result of a `GROUP BY` would.

What this buys over N0. A schema says a quantity is between one and nine; it
does not say that most orders are for one item and almost none for nine. Data
drawn uniformly inside the bounds gives every rare case the same weight as the
common one, so a cache hit rate measured on it means nothing, a page laid out
for it looks nothing like the real one, and the slow query stays fast.

What is sampled here is the marginal distribution of each column, one column at
a time: the observed count per category, and the observed count per bucket of a
numeric column. Nothing is fitted and nothing is learnt — the observed table is
the model. Which also means the joint distribution is lost, and the tests next
to this file show what that costs.

Determinism works exactly as in N0, and for the same reason: a data set that is
only reproducible on average is not reproducible.
"""

FNV_OFFSET = 2166136261
FNV_PRIME = 16777619
MASK32 = 0xFFFFFFFF

UNIT = "\x1f"


def stable_hash(text: str) -> int:
    """FNV-1a on 32 bits, identical to the JavaScript version of this file."""
    digest = FNV_OFFSET
    for char in text:
        digest = ((digest ^ ord(char)) * FNV_PRIME) & MASK32
    return digest


def draw(seed: str, field: str, row: int) -> int:
    """One 32-bit integer per cell, independent of the other cells."""
    return stable_hash(f"{seed}{UNIT}{field}{UNIT}{row}")


def pick(counts: list[int], number: int) -> int:
    """
    Index of the bucket a drawn integer falls into, in proportion to `counts`.

    Cumulating integers rather than normalising to probabilities keeps the
    result exact, and identical in both languages: no floating point is
    involved anywhere in the decision.
    """
    total = sum(counts)
    if total <= 0:
        raise ValueError("a distribution needs at least one observation")
    target = number % total
    for index, count in enumerate(counts):
        if target < count:
            return index
        target -= count
    raise AssertionError("unreachable: the target is below the total")


def _value(spec: dict, seed: str, field: str, row: int):
    kind = spec["type"]
    if kind == "categorical":
        # `counts` is a label to observed-count mapping, straight out of a
        # `GROUP BY`. A label seen zero times is never drawn, which is the
        # honest behaviour: it did not happen.
        #
        # The labels are sorted rather than taken in the order the query
        # returned them, so the same observed table always gives the same data.
        # It is also what keeps the two languages together: a JavaScript object
        # reorders its numeric-looking keys, and a postcode is one of those.
        labels = sorted(spec["counts"])
        return labels[pick([spec["counts"][label] for label in labels], draw(seed, field, row))]
    if kind == "histogram":
        edges, counts = spec["edges"], spec["counts"]
        if len(edges) != len(counts) + 1:
            raise ValueError("a histogram needs one more edge than it has buckets")
        bucket = pick(counts, draw(seed, field, row))
        low, high = edges[bucket], edges[bucket + 1]
        # A second, independent draw places the value inside its bucket. Within
        # a bucket the shape is unknown, so uniform is the only honest choice.
        return low + draw(seed, f"{field}{UNIT}within", row) % max(high - low, 1)
    raise ValueError(f"unknown distribution type {kind!r}")


def sample_rows(distributions: dict, count: int, seed: str) -> list[dict]:
    """
    Return `count` rows, each column sampled from its own observed distribution.

    `distributions` maps a field name to either
    {"type": "categorical", "counts": {label: observed count}} or
    {"type": "histogram", "edges": [...], "counts": [...]}, where the buckets
    are half-open and `edges` holds one more value than `counts`.
    """
    return [
        {field: _value(spec, seed, field, row) for field, spec in distributions.items()}
        for row in range(count)
    ]

/**
 * Sample a test data set from the distributions observed in production.
 *
 * Rung N1. No dependency, no file written, and still no network: the
 * distributions come from the caller, as the result of a `GROUP BY` would.
 *
 * What this buys over N0. A schema says a quantity is between one and nine; it
 * does not say that most orders are for one item and almost none for nine.
 * Data drawn uniformly inside the bounds gives every rare case the same weight
 * as the common one, so a cache hit rate measured on it means nothing, a page
 * laid out for it looks nothing like the real one, and the slow query stays
 * fast.
 *
 * What is sampled here is the marginal distribution of each column, one column
 * at a time: the observed count per category, and the observed count per
 * bucket of a numeric column. Nothing is fitted and nothing is learnt — the
 * observed table is the model. Which also means the joint distribution is
 * lost, and the tests next to this file show what that costs.
 *
 * Determinism works exactly as in N0, and for the same reason: a data set that
 * is only reproducible on average is not reproducible.
 */

const FNV_OFFSET = 2166136261;
const FNV_PRIME = 16777619;

// Unit separator: it joins the parts of a cell key, and appears in none of them.
const UNIT = '\u001f';

/** FNV-1a on 32 bits, identical to the Python version of this file. */
export function stableHash(text) {
  let digest = FNV_OFFSET;
  for (const char of text) {
    digest = Math.imul(digest ^ char.codePointAt(0), FNV_PRIME) >>> 0;
  }
  return digest;
}

/** One 32-bit integer per cell, independent of the other cells. */
export function draw(seed, field, row) {
  return stableHash(`${seed}${UNIT}${field}${UNIT}${row}`);
}

/**
 * Index of the bucket a drawn integer falls into, in proportion to `counts`.
 *
 * Cumulating integers rather than normalising to probabilities keeps the
 * result exact, and identical in both languages: no floating point is involved
 * anywhere in the decision.
 */
export function pick(counts, number) {
  const total = counts.reduce((sum, count) => sum + count, 0);
  if (total <= 0) throw new RangeError('a distribution needs at least one observation');
  let target = number % total;
  for (let index = 0; index < counts.length; index += 1) {
    if (target < counts[index]) return index;
    target -= counts[index];
  }
  throw new Error('unreachable: the target is below the total');
}

function value(spec, seed, field, row) {
  switch (spec.type) {
    case 'categorical': {
      // `counts` maps a label to its observed count, straight out of a
      // `GROUP BY`. A label seen zero times is never drawn, which is the
      // honest behaviour: it did not happen.
      //
      // The labels are sorted rather than taken in the order the query
      // returned them, so the same observed table always gives the same data.
      // It is also what keeps the two languages together: an object reorders
      // its numeric-looking keys, and a postcode is one of those.
      const labels = Object.keys(spec.counts).sort();
      return labels[pick(labels.map((label) => spec.counts[label]), draw(seed, field, row))];
    }
    case 'histogram': {
      const { edges, counts } = spec;
      if (edges.length !== counts.length + 1) {
        throw new RangeError('a histogram needs one more edge than it has buckets');
      }
      const bucket = pick(counts, draw(seed, field, row));
      const [low, high] = [edges[bucket], edges[bucket + 1]];
      // A second, independent draw places the value inside its bucket. Within
      // a bucket the shape is unknown, so uniform is the only honest choice.
      return low + (draw(seed, `${field}${UNIT}within`, row) % Math.max(high - low, 1));
    }
    default:
      throw new RangeError(`unknown distribution type ${JSON.stringify(spec.type)}`);
  }
}

/**
 * Return `count` rows, each column sampled from its own observed distribution.
 *
 * `distributions` maps a field name to either
 * {type: 'categorical', counts: {label: observed count}} or
 * {type: 'histogram', edges, counts}, where the buckets are half-open and
 * `edges` holds one more value than `counts`.
 */
export function sampleRows(distributions, count, seed) {
  const rows = [];
  for (let row = 0; row < count; row += 1) {
    const built = {};
    for (const [field, spec] of Object.entries(distributions)) {
      built[field] = value(spec, seed, field, row);
    }
    rows.push(built);
  }
  return rows;
}

"""
Write the text fields of a test data set by asking a general-purpose model.

Rung N3. This is the option people reach for first, and it does buy something
real: a support ticket that reads like a support ticket, with the typos, the
capitals and the two questions in one sentence that no template produces.

It also gives up the property the lower rungs were built on. There is no seed
here. Two calls with the same prompt return different rows, so the data set has
to be generated once and then stored, like a fixture, not rebuilt on demand.

And the model guarantees nothing: not the keys you asked for, not the number of
rows, not the uniqueness of an identifier. Look at how much of this file is
checking rather than asking. That plumbing is the real cost of the rung, and it
is the part the tests can cover, because the model itself is not testable.
"""

from __future__ import annotations

import json

MAX_ROWS = 50  # beyond that the answer comes back truncated more often than not


class GenerationUnavailable(Exception):
    """The provider could not be reached, or never returned a usable data set."""


def build_prompt(fields: list[str], count: int, unique_field: str | None) -> str:
    """The instructions, kept next to the checks that verify they were followed."""
    lines = [
        f"Write {count} rows of test data for a fictional application.",
        "Each row is a JSON object with exactly these keys, and string values:",
        ", ".join(fields) + ".",
        "Invent every value: it must match no real person, company or address.",
        f"Answer with JSON only: a list of {count} objects, and nothing else.",
    ]
    if unique_field:
        lines.insert(3, f"Every value of `{unique_field}` must differ from the others.")
    return "\n".join(lines)


def check(rows, fields: list[str], count: int, unique_field: str | None) -> None:
    """
    Verify what the model was asked for. Nothing here is redundant.

    Each of these failures is one this rung produces in practice: a row short,
    a key renamed to its plural, an empty string, the same name twice.
    """
    if not isinstance(rows, list) or len(rows) != count:
        raise ValueError(f"expected a list of {count} rows")
    for row in rows:
        if not isinstance(row, dict) or set(row) != set(fields):
            raise ValueError(f"a row does not carry exactly the keys {fields}")
        if not all(isinstance(value, str) and value.strip() for value in row.values()):
            raise ValueError("a value is empty, or is not a string")
    if unique_field:
        values = [row[unique_field] for row in rows]
        if len(set(values)) != len(values):
            raise ValueError(f"the model repeated a value of {unique_field!r}")


def write_rows(
    fields: list[str],
    count: int,
    *,
    unique_field: str | None = None,
    client=None,
    attempts: int = 3,
    temperature: float = 1.0,
) -> list[dict]:
    """
    Return `count` rows of invented text, or raise rather than return junk.

    `client` is injected so this function can be tested without a network call.
    In production it defaults to a real provider client.

    The temperature is high on purpose: varied prose is the only reason to be
    on this rung at all. It is also why the answer has to be checked, and why
    the same call twice gives two different data sets.
    """
    if client is None:  # pragma: no cover - needs a key and a network
        from openai import OpenAI

        client = OpenAI()

    # A model charges by the token. Refusing an oversized batch before calling
    # is not an optimisation, it is a cost control.
    if not 0 < count <= MAX_ROWS:
        raise ValueError(f"ask for between one and {MAX_ROWS} rows at a time")

    prompt = build_prompt(fields, count, unique_field)
    last_error: Exception | None = None
    for _ in range(attempts):
        try:
            answer = client.complete(prompt=prompt, temperature=temperature)
            rows = json.loads(answer)
            check(rows, fields, count, unique_field)
            return rows
        except Exception as error:  # noqa: BLE001 - a bad answer is retried like a failure
            last_error = error
    raise GenerationUnavailable(str(last_error))

/**
 * Write the text fields of a test data set by asking a general-purpose model.
 *
 * Rung N3. This is the option people reach for first, and it does buy
 * something real: a support ticket that reads like a support ticket, with the
 * typos, the capitals and the two questions in one sentence that no template
 * produces.
 *
 * It also gives up the property the lower rungs were built on. There is no
 * seed here. Two calls with the same prompt return different rows, so the data
 * set has to be generated once and then stored, like a fixture, not rebuilt on
 * demand.
 *
 * And the model guarantees nothing: not the keys you asked for, not the number
 * of rows, not the uniqueness of an identifier. Look at how much of this file
 * is checking rather than asking. That plumbing is the real cost of the rung,
 * and it is the part the tests can cover, because the model itself is not
 * testable.
 */

// Beyond that the answer comes back truncated more often than not.
export const MAX_ROWS = 50;

export class GenerationUnavailable extends Error {}

/** The instructions, kept next to the checks that verify they were followed. */
export function buildPrompt(fields, count, uniqueField) {
  const lines = [
    `Write ${count} rows of test data for a fictional application.`,
    'Each row is a JSON object with exactly these keys, and string values:',
    `${fields.join(', ')}.`,
    'Invent every value: it must match no real person, company or address.',
    `Answer with JSON only: a list of ${count} objects, and nothing else.`,
  ];
  if (uniqueField) {
    lines.splice(3, 0, `Every value of \`${uniqueField}\` must differ from the others.`);
  }
  return lines.join('\n');
}

/**
 * Verify what the model was asked for. Nothing here is redundant.
 *
 * Each of these failures is one this rung produces in practice: a row short, a
 * key renamed to its plural, an empty string, the same name twice.
 */
export function check(rows, fields, count, uniqueField) {
  if (!Array.isArray(rows) || rows.length !== count) {
    throw new TypeError(`expected a list of ${count} rows`);
  }
  for (const row of rows) {
    const keys = row && typeof row === 'object' ? Object.keys(row) : [];
    if (keys.length !== fields.length || !fields.every((field) => keys.includes(field))) {
      throw new TypeError(`a row does not carry exactly the keys ${fields.join(', ')}`);
    }
    if (!Object.values(row).every((value) => typeof value === 'string' && value.trim())) {
      throw new TypeError('a value is empty, or is not a string');
    }
  }
  if (uniqueField) {
    const values = rows.map((row) => row[uniqueField]);
    if (new Set(values).size !== values.length) {
      throw new TypeError(`the model repeated a value of ${uniqueField}`);
    }
  }
}

/**
 * Return `count` rows of invented text, or throw rather than return junk.
 *
 * The temperature is high on purpose: varied prose is the only reason to be on
 * this rung at all. It is also why the answer has to be checked, and why the
 * same call twice gives two different data sets.
 *
 * @param {string[]} fields
 * @param {number} count
 * @param {object} options
 * @param {string} [options.uniqueField] field whose values must not repeat
 * @param {{complete: Function}} [options.client] injected so this can be
 *   tested without a network call; defaults to a real provider client
 * @param {number} [options.attempts]
 * @param {number} [options.temperature]
 */
export async function writeRows(fields, count, options = {}) {
  const { uniqueField, attempts = 3, temperature = 1 } = options;
  let { client } = options;
  if (!client) {
    // Needs a key and a network, so it is never reached in the tests.
    const { OpenAI } = await import('openai');
    client = new OpenAI();
  }

  // A model charges by the token. Refusing an oversized batch before calling
  // is not an optimisation, it is a cost control.
  if (!(count > 0 && count <= MAX_ROWS)) {
    throw new RangeError(`ask for between one and ${MAX_ROWS} rows at a time`);
  }

  const prompt = buildPrompt(fields, count, uniqueField);
  let lastError;
  for (let attempt = 0; attempt < attempts; attempt += 1) {
    try {
      const rows = JSON.parse(await client.complete({ prompt, temperature }));
      check(rows, fields, count, uniqueField);
      return rows;
    } catch (error) {
      // A bad answer is retried exactly like a provider failure.
      lastError = error;
    }
  }
  throw new GenerationUnavailable(String(lastError));
}