"""Tests for the DynamoDB connector."""

from typing import Any

import boto3
import pytest_mock
from moto import mock_aws

from assets.connectors import dynamodb

TEST_TABLE_NAME = "test-table"
TABLE_PKEY = "order_id"


def get_mock_orders(num_items: int = 1) -> list[dict[str, str]]:
    """Return fake order dynamo data."""
    return [
        {
            "order_id": str(i),
            "encoder_id": f"value_{i}",
        }
        for i in range(1, num_items + 1)
    ]


def mock_ddb_table(client: Any, items: list[dict[str, str]]) -> None:
    """Mock DDB Table."""
    client.create_table(
        TableName=TEST_TABLE_NAME,
        KeySchema=[{"AttributeName": TABLE_PKEY, "KeyType": "HASH"}],
        AttributeDefinitions=[{"AttributeName": TABLE_PKEY, "AttributeType": "S"}],
        ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1},
    )
    for item in items:
        client.put_item(
            TableName=TEST_TABLE_NAME,
            Item={
                "order_id": {"S": item[TABLE_PKEY]},
                "encoder_id": {"S": item["encoder_id"]},
            },
        )


def test_get_dynamodb_items(mocker: pytest_mock.MockerFixture) -> None:
    """Test retrieving DynamoDB items in batches."""
    mock_items = get_mock_orders(3)
    keys = [str(v[TABLE_PKEY]) for v in mock_items]
    with mock_aws():
        mocked_sleep = mocker.patch("assets.connectors.dynamodb.time.sleep")
        mocked_get_unprocessed_keys = mocker.patch(
            "assets.connectors.dynamodb.DDBClient.get_unprocessed_keys"
        )
        # Pretend we've received non-empty UnprocessedKeys on the first call.
        mocked_get_unprocessed_keys.side_effect = [["3"], []]

        client = boto3.client("dynamodb")
        mock_ddb_table(client, mock_items)
        result = dynamodb.ddb_connector.batch_get_items(
            TEST_TABLE_NAME, TABLE_PKEY, keys, "order_id,encoder_id"
        )
    assert result == [
        {"encoder_id": "value_1", TABLE_PKEY: "1"},
        {"encoder_id": "value_2", TABLE_PKEY: "2"},
        {"encoder_id": "value_3", TABLE_PKEY: "3"},
        {"encoder_id": "value_3", TABLE_PKEY: "3"},
    ]
    mocked_sleep.call_count = 1
    mocked_sleep.assert_called_with(0.3)


def test_get_unprocessed_keys() -> None:
    """Test get_unprocessed_keys."""
    actual = dynamodb.ddb_connector.get_unprocessed_keys({}, "foo", "bar")
    assert actual == []
    actual = dynamodb.ddb_connector.get_unprocessed_keys(
        {"UnprocessedKeys": {}}, "foo", "bar"
    )
    assert actual == []
    actual = dynamodb.ddb_connector.get_unprocessed_keys(
        {"UnprocessedKeys": {"foo": None}}, "foo", "bar"
    )
    assert actual == []
    actual = dynamodb.ddb_connector.get_unprocessed_keys(
        {"UnprocessedKeys": {"foo": {"Keys": None}}}, "foo", "bar"
    )
    assert actual == []
    actual = dynamodb.ddb_connector.get_unprocessed_keys(
        {
            "UnprocessedKeys": {
                "foo": {
                    "Keys": [
                        {"bar": {"S": "1"}},
                        {"bar": {"S": "2"}},
                        {"bar": {"S": "3"}},
                    ]
                }
            }
        },
        "foo",
        "bar",
    )
    assert actual == ["1", "2", "3"]