test_query_api.py

"""Tests for the Query API."""

import datetime
import json
from unittest.mock import MagicMock

import boto3
import pytest

from sds_data_manager.lambda_code.SDSCode.api_lambdas import query_api
from sds_data_manager.lambda_code.SDSCode.database import models


def param_not_valid_in_response(response_body, param, table):
    """Check if error message contains expected content for invalid parameter."""
    error_msg = f"{param} is not a valid query parameter for {table} table"
    return error_msg in response_body


def _populate_test_data(session):
    """Put a filepath into the test data."""
    filepath = "test/file/path/imap_hit_l0_raw_20251107_v001.pkts"

    metadata_params = {
        "file_path": filepath,
        "instrument": "hit",
        "data_level": "l0",
        "descriptor": "raw",
        "start_date": datetime.datetime.strptime("20251107", "%Y%m%d"),
        "version": "v001",
        "extension": "pkts",
        "ingestion_date": datetime.datetime.strptime(
            "2025-11-07 10:13:12+00:00", "%Y-%m-%d %H:%M:%S%z"
        ),
        "released": True,
    }

    # Add data to the ScienceFiles table and return the session
    session.add(models.ScienceFiles(**metadata_params))
    session.commit()


@pytest.fixture(autouse=True)
def mock_head_object(monkeypatch):
    """Patch boto3's S3 client head_object to always return success."""
    # Create a mock that always returns a success response
    mock_head_object = MagicMock(
        return_value={"ResponseMetadata": {"HTTPStatusCode": 200}}
    )

    # Also patch the boto3.resource to handle the bucket name issue
    mock_bucket = MagicMock()
    mock_bucket.name = "mock-bucket-name"  # Add a name attribute that can be used

    # Patch both s3_client.head_object and the bucket name access
    def mock_client(*args, **kwargs):
        mock = MagicMock()
        mock.head_object = mock_head_object
        return mock

    def mock_resource(*args, **kwargs):
        mock = MagicMock()
        mock.Bucket.return_value = mock_bucket
        return mock

    monkeypatch.setattr(boto3, "client", mock_client)
    monkeypatch.setattr(boto3, "resource", mock_resource)


@pytest.fixture
def expected_response():
    """Return the expected response."""
    expected_response = json.dumps(
        [
            {
                "file_path": "test/file/path/imap_hit_l0_raw_20251107_v001.pkts",
                "instrument": "hit",
                "data_level": "l0",
                "descriptor": "raw",
                "start_date": "20251107",
                "repointing": None,
                "version": "v001",
                "extension": "pkts",
                "ingestion_date": "20251107 10:13:12",
                "cr": None,
                "crid": None,
                "released": True,
            }
        ]
    )
    return expected_response


def test_query_result_body(session):
    """Tests that the query result body can be loaded."""
    _populate_test_data(session)
    event = {"queryStringParameters": {}}

    returned_query = query_api.lambda_handler(event=event, context={})

    assert json.loads(returned_query["body"])


def test_query_result_header(session):
    """Tests that the query result header is json."""
    _populate_test_data(session)
    event = {"queryStringParameters": {}}

    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["headers"] is not None
    assert returned_query["headers"]["Content-Type"] == "application/json"


def test_start_date_query(session, expected_response):
    """Test that start date can be queried."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"start_date": "20251101"}}

    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    # Parse both JSON objects and compare the data rather than the string representation
    assert json.loads(returned_query["body"]) == json.loads(expected_response)


def test_end_date_query(session, expected_response):
    """Test that end date can be queried."""
    _populate_test_data(session)
    event = {
        "queryStringParameters": {"start_date": "20251101"},
    }
    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    # Parse both JSON objects and compare the data rather than the string representation
    assert json.loads(returned_query["body"]) == json.loads(expected_response)


def test_start_and_end_date_query(session, expected_response):
    """Test that both start and end date can be queried."""
    event = {
        "queryStringParameters": {"start_date": "20251101", "end_date": "20251201"}
    }
    _populate_test_data(session)
    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    # Parse both JSON objects and compare the data rather than the string representation
    assert json.loads(returned_query["body"]) == json.loads(expected_response)


def test_empty_start_date_query(session):
    """Test that a start_date query with no matches returns an empty list."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"start_date": "20261101"}}
    expected_response = json.dumps([])
    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    assert returned_query["body"] == expected_response


def test_empty_end_date_query(session):
    """Test that an end_date query with no matches returns an empty list."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"start_date": "20261101"}}
    expected_response = json.dumps([])
    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    assert returned_query["body"] == expected_response


def test_non_date_query(session, expected_response):
    """Test that a non-date parameters can be queried."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"instrument": "hit"}}

    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    # Parse both JSON objects and compare the data rather than the string representation
    assert json.loads(returned_query["body"]) == json.loads(expected_response)


def test_ingestion_start_date_query(session, expected_response):
    """Test that ingestion_start_date can be queried."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"ingestion_start_date": "20251107"}}

    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    # Parse both JSON objects and compare the data rather than the string representation
    assert json.loads(returned_query["body"]) == json.loads(expected_response)


def test_ingestion_end_date_query(session, expected_response):
    """Test that ingestion_end_date can be queried."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"ingestion_end_date": "20251107"}}

    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    # Parse both JSON objects and compare the data rather than the string representation
    assert json.loads(returned_query["body"]) == json.loads(expected_response)


def test_ingestion_start_and_end_date_query(session, expected_response):
    """Test that both ingestion_start_date and ingestion_end_date can be queried."""
    _populate_test_data(session)
    event = {
        "queryStringParameters": {
            "ingestion_start_date": "20251106",
            "ingestion_end_date": "20251108",
        }
    }

    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    # Parse both JSON objects and compare the data rather than the string representation
    assert json.loads(returned_query["body"]) == json.loads(expected_response)


def test_empty_ingestion_start_date_query(session):
    """Test that an ingestion_start_date query with no matches returns an empty list."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"ingestion_start_date": "20261101"}}
    expected_response = json.dumps([])
    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    assert returned_query["body"] == expected_response


def test_empty_ingestion_end_date_query(session):
    """Test that an ingestion_end_date query with no matches returns an empty list."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"ingestion_end_date": "20251101"}}
    expected_response = json.dumps([])
    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    assert returned_query["body"] == expected_response


def test_empty_ingestion_start_and_end_date_query(session):
    """Test that ingestion params with no matches returns an empty list."""
    _populate_test_data(session)
    event = {
        "queryStringParameters": {
            "ingestion_start_date": "20261101",
            "ingestion_end_date": "20261110",
        }
    }
    expected_response = json.dumps([])
    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    assert returned_query["body"] == expected_response


def test_empty_non_date_query(session):
    """Test that a non-date query with no matches returns an empty list."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"data_level": "l2"}}
    expected_response = json.dumps([])
    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    assert returned_query["body"] == expected_response


def test_multi_param_query(session, expected_response):
    """Test that multiple parameters can be queried."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"instrument": "hit", "data_level": "l0"}}

    returned_query = query_api.lambda_handler(event=event, context={})
    assert returned_query["statusCode"] == 200
    # Parse both JSON objects and compare the data rather than the string representation
    assert json.loads(returned_query["body"]) == json.loads(expected_response)


def test_invalid_query(session):
    """Test that invalid parameters return a 400 status with explanation."""
    _populate_test_data(session)
    event = {"queryStringParameters": {"size": "500"}}
    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 400
    # Check if error message contains the expected content
    assert param_not_valid_in_response(returned_query["body"], "size", "science")


def _populate_test_data_ancillary_table(session):
    """Put a filepath into the test data for the ancillary table."""
    filepath = "test/ancillary/file/path/imap_mag_test_20210101_v001.csv"

    metadata_params = {
        "file_path": filepath,
        "instrument": "mag",
        "descriptor": "test",
        "start_date": datetime.datetime.strptime("20210101", "%Y%m%d"),
        "version": "v001",
        "extension": "csv",
        "ingestion_date": datetime.datetime.strptime(
            "2021-01-01 10:13:12+00:00", "%Y-%m-%d %H:%M:%S%z"
        ),
        "released": True,
    }

    # Add data to the AncillaryFiles table and return the session
    session.add(models.AncillaryFiles(**metadata_params))
    session.commit()


@pytest.fixture
def expected_response_ancillary_table():
    """Return the expected response for ancillary table."""
    expected_response = json.dumps(
        [
            {
                "file_path": "test/ancillary/file/path/imap_mag_test_20210101_v001.csv",
                "instrument": "mag",
                "descriptor": "test",
                "start_date": "20210101",
                "end_date": None,
                "version": "v001",
                "extension": "csv",
                "ingestion_date": "20210101 10:13:12",
                "released": True,
            }
        ]
    )
    return expected_response


def test_query_result_body_ancillary_table(session):
    """Tests that the query result body can be loaded for ancillary table."""
    _populate_test_data_ancillary_table(session)
    event = {"queryStringParameters": {"table": "ancillary"}}

    returned_query = query_api.lambda_handler(event=event, context={})

    assert json.loads(returned_query["body"])


def test_query_ancillary_table(session, expected_response_ancillary_table):
    """Test querying the ancillary table with a valid parameter."""
    _populate_test_data_ancillary_table(session)

    event = {"queryStringParameters": {"instrument": "mag", "table": "ancillary"}}

    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 200
    # Parse both JSON objects and compare the data rather than the string representation
    assert json.loads(returned_query["body"]) == json.loads(
        expected_response_ancillary_table
    )


def test_invalid_param_ancillary_query(session):
    """Test invalid parameter on the ancillary table."""
    _populate_test_data_ancillary_table(session)

    event = {"queryStringParameters": {"repointing": "123", "table": "ancillary"}}

    returned_query = query_api.lambda_handler(event=event, context={})

    assert returned_query["statusCode"] == 400
    # Check if error message contains the expected content
    assert param_not_valid_in_response(
        returned_query["body"], "repointing", "ancillary"
    )