test_list_axis_reader.py

"""Tests for list-axis scatter-gather readers (issue #23 follow-on)."""

from __future__ import annotations

import pytest

from security_scanner.storage.adapters.nosql_db.list_axis import (
    REPO_LIST_AXIS,
    TARGET_LIST_AXIS,
    legacy_list_axis_pk,
    list_axis_projection_for_item,
)
from security_scanner.storage.adapters.nosql_db.list_axis_reader import (
    read_list_axis,
    read_list_axis_ordered,
)


class _FakeGsiTable:
    def __init__(self) -> None:
        self.items: list[dict] = []
        self.queried_pks: list[str] = []

    def add(self, item: dict) -> None:
        self.items.append(item)

    def query(self, **kwargs) -> dict:
        values = kwargs["ExpressionAttributeValues"]
        pk = values[":pk"]
        self.queried_pks.append(pk)
        prefix = values.get(":p")
        matched = [
            i
            for i in self.items
            if i.get("gsi1pk") == pk
            and (prefix is None or str(i.get("gsi1sk", "")).startswith(prefix))
        ]
        matched.sort(
            key=lambda i: i.get("gsi1sk", ""),
            reverse=not kwargs.get("ScanIndexForward", True),
        )
        if "Limit" in kwargs:
            matched = matched[: kwargs["Limit"]]
        return {"Items": matched}


def _target_item(name: str) -> dict:
    item = {
        "PK": f"SCAN_TARGET#https://x/{name}",
        "SK": "META",
        "entityType": "SCAN_TARGET",
        "name": name,
        "url": f"https://x/{name}",
    }
    item.update(list_axis_projection_for_item(TARGET_LIST_AXIS, item))
    return item


def _repo_item(repo_key: str, updated_at: str, *, sharded: bool = True) -> dict:
    item = {
        "PK": f"REPO#{repo_key}",
        "SK": "META",
        "entityType": "REPO_META",
        "repoKey": repo_key,
        "updatedAt": updated_at,
    }
    if sharded:
        item.update(list_axis_projection_for_item(REPO_LIST_AXIS, item))
    else:  # legacy unsharded row, no version attr
        item["gsi1pk"] = legacy_list_axis_pk("REPO_LIST#ALL")
        item["gsi1sk"] = f"UPDATED#{updated_at}#{repo_key}"
    return item


def test_flat_read_fans_out_all_shards_and_merges():
    table = _FakeGsiTable()
    for n in range(6):
        table.add(_target_item(f"org/repo{n}"))

    result = read_list_axis(
        table,
        spec=TARGET_LIST_AXIS,
        partition_root="TARGET_LIST#ALL",
        gsi_sk_prefix="TARGET#",
    )

    assert len(table.queried_pks) == TARGET_LIST_AXIS.shard_count  # 8, no legacy
    assert [i["gsi1sk"] for i in result] == [f"TARGET#org/repo{n}" for n in range(6)]


def test_flat_read_default_issues_no_legacy_query():
    table = _FakeGsiTable()
    read_list_axis(
        table,
        spec=TARGET_LIST_AXIS,
        partition_root="TARGET_LIST#ALL",
        gsi_sk_prefix="TARGET#",
    )
    assert legacy_list_axis_pk("TARGET_LIST#ALL") not in table.queried_pks


def test_flat_read_include_legacy_merges_and_dedupes_prefer_higher_version():
    table = _FakeGsiTable()
    table.add(_repo_item("org/sharded", "2026-06-19T00:00:00+00:00", sharded=True))
    # legacy-only row
    table.add(_repo_item("org/legacy", "2026-06-18T00:00:00+00:00", sharded=False))
    # duplicate (PK,SK): both legacy and sharded for org/dup
    dup_sharded = _repo_item("org/dup", "2026-06-17T00:00:00+00:00", sharded=True)
    dup_legacy = _repo_item("org/dup", "2026-06-17T00:00:00+00:00", sharded=False)
    table.add(dup_sharded)
    table.add(dup_legacy)

    result = read_list_axis(
        table,
        spec=REPO_LIST_AXIS,
        partition_root="REPO_LIST#ALL",
        gsi_sk_prefix="UPDATED#",
        include_legacy=True,
    )

    keys = [(i["PK"], i["SK"]) for i in result]
    assert ("REPO#org/legacy", "META") in keys
    assert keys.count(("REPO#org/dup", "META")) == 1  # deduped
    dup = next(i for i in result if i["PK"] == "REPO#org/dup")
    assert dup.get("listAxisVersion") == 1  # higher-version (sharded) wins


def test_flat_read_fails_closed_on_shard_error():
    class _Boom(_FakeGsiTable):
        def query(self, **kwargs):
            # shard_count=8 → single-digit buckets "0".."7"
            if kwargs["ExpressionAttributeValues"][":pk"].endswith("#SHARD#2"):
                raise RuntimeError("shard 02 down")
            return super().query(**kwargs)

    table = _Boom()
    table.add(_target_item("org/repo"))
    with pytest.raises(RuntimeError, match="shard 02 down"):
        read_list_axis(
            table,
            spec=TARGET_LIST_AXIS,
            partition_root="TARGET_LIST#ALL",
            gsi_sk_prefix="TARGET#",
        )


def test_ordered_read_returns_global_newest_first_within_limit():
    table = _FakeGsiTable()
    # 10 repos with increasing timestamps spread across shards
    for n in range(10):
        ts = f"2026-06-{10 + n:02d}T00:00:00+00:00"
        table.add(_repo_item(f"org/repo{n:02d}", ts))

    result = read_list_axis_ordered(
        table,
        spec=REPO_LIST_AXIS,
        partition_root="REPO_LIST#ALL",
        gsi_sk_prefix="UPDATED#",
        limit=3,
    )

    # newest 3 are repo09, repo08, repo07
    assert [i["repoKey"] for i in result] == ["org/repo09", "org/repo08", "org/repo07"]


def test_ordered_read_limit_none_returns_all_sorted():
    table = _FakeGsiTable()
    for n in range(5):
        table.add(_repo_item(f"org/repo{n}", f"2026-06-1{n}T00:00:00+00:00"))
    result = read_list_axis_ordered(
        table,
        spec=REPO_LIST_AXIS,
        partition_root="REPO_LIST#ALL",
        gsi_sk_prefix="UPDATED#",
        limit=None,
    )
    assert [i["repoKey"] for i in result] == [f"org/repo{n}" for n in (4, 3, 2, 1, 0)]


def test_ordered_read_limit_exceeds_total():
    table = _FakeGsiTable()
    table.add(_repo_item("org/a", "2026-06-10T00:00:00+00:00"))
    result = read_list_axis_ordered(
        table,
        spec=REPO_LIST_AXIS,
        partition_root="REPO_LIST#ALL",
        gsi_sk_prefix="UPDATED#",
        limit=50,
    )
    assert [i["repoKey"] for i in result] == ["org/a"]


def test_ordered_read_default_issues_no_legacy_query():
    table = _FakeGsiTable()
    for n in range(3):
        table.add(_repo_item(f"org/r{n}", f"2026-06-1{n}T00:00:00+00:00"))

    read_list_axis_ordered(
        table,
        spec=REPO_LIST_AXIS,
        partition_root="REPO_LIST#ALL",
        gsi_sk_prefix="UPDATED#",
        limit=2,
    )

    assert legacy_list_axis_pk("REPO_LIST#ALL") not in table.queried_pks
    assert len(table.queried_pks) == REPO_LIST_AXIS.shard_count


def test_ordered_read_include_legacy_dedupes_without_truncating_window():
    table = _FakeGsiTable()
    # legacy + sharded duplicate of the SAME newest row, plus older distinct rows
    newest = "2026-06-20T00:00:00+00:00"
    table.add(_repo_item("org/newest", newest, sharded=True))
    table.add(_repo_item("org/newest", newest, sharded=False))  # legacy dup
    table.add(_repo_item("org/old1", "2026-06-11T00:00:00+00:00", sharded=False))
    table.add(_repo_item("org/old2", "2026-06-10T00:00:00+00:00", sharded=False))

    result = read_list_axis_ordered(
        table,
        spec=REPO_LIST_AXIS,
        partition_root="REPO_LIST#ALL",
        gsi_sk_prefix="UPDATED#",
        limit=2,
        include_legacy=True,
    )

    # distinct top-2: newest (deduped to one), then old1 — not silently short
    assert [i["repoKey"] for i in result] == ["org/newest", "org/old1"]
    assert legacy_list_axis_pk("REPO_LIST#ALL") in table.queried_pks