from io import BytesIO

import pandas
import pytest

from common.src.aws import new_s3_client
from .....src import config
from .....src.logic.consumer import handler
from common.src.enums import AWSPayload

# Test for the consumer handler, using a real label and period ID,
# which we know will produce results.
TEST_LABEL_ID: int = 24483
TEST_PERIOD_ID: int = 315

PREFIX: str = "__test/"


@pytest.fixture(autouse=True)
def patcher_generate_key(monkeypatch):
    """Patcher to generate a test specific S3 key, by prefixing the
    S3 key with a test prefix. This is to avoid overwriting real
    data in S3 and have it placed in a test folder.
    """
    orig = handler._generate_key
    monkeypatch.setattr(
        handler,
        handler._generate_key.__name__,
        lambda *a, **kw: f"{PREFIX}{orig(*a, **kw)}",
    )


@pytest.fixture
def fetch_from_s3(request):
    """Fetches the S3 object from the given bucket and key, with cleanup
    after the test.
    """
    to_cleanup = []

    def _fetch(bucket: str, key: str):
        s3_client = new_s3_client(config.AWS_REGION)
        response = s3_client.get_object(Bucket=bucket, Key=key)
        to_cleanup.append((s3_client, bucket, key))
        return BytesIO(response[AWSPayload.BODY.value].read())

    def cleanup():
        for s3_client, bucket, key in to_cleanup:
            s3_client.delete_object(Bucket=bucket, Key=key)
            response = s3_client.list_objects_v2(Bucket=bucket, Prefix=key)
            assert "Contents" not in response, f"S3 object {key} was not deleted"

    request.addfinalizer(cleanup)
    return _fetch


def test_run(fetch_from_s3):
    bucket, key = handler.run(TEST_LABEL_ID, TEST_PERIOD_ID)
    s3_blob = fetch_from_s3(bucket, key)
    sheets = pandas.read_excel(s3_blob, sheet_name=None)
    assert len(sheets) >= 1