"""Unit tests for Amazon DataPulse exec."""

import argparse

import pytest

from feed_ingestion.flows.amazon_datapulse import config
from feed_ingestion.flows.amazon_datapulse import (  # noqa: A004
    exec as flow_exec_module,
)
from tests.test_exec_flow import FlowExecTestBase


class TestFlowExec(FlowExecTestBase):
    """Tests for FlowExec parser configuration."""

    @pytest.fixture(autouse=True)
    def setup(self):
        """Set up FlowExec instance."""
        self.flow_exec = flow_exec_module.FlowExec()

    PARSER_ARGS = FlowExecTestBase.PARSER_ARGS | {
        'days',
        'skip',
        'report',
        'partition',
    }

    ARGS_REQUIRED = FlowExecTestBase.ARGS_REQUIRED | {
        'report',
    }

    AGRS_CHOICES = FlowExecTestBase.AGRS_CHOICES | {
        'report': ['ALL'] + list(config.reports.keys()),
    }

    def test_generate_contexts(self):
        """Test generate_contexts with a single country and report."""
        self.flow_exec.args = argparse.Namespace(
            country='US', report='fraud_report'
        )
        result = list(self.flow_exec.generate_contexts())
        assert result == [
            {
                'report_name': 'fraud_report',
            }
        ]

    def test_generate_contexts_all_reports(self):
        """Test generate_contexts expands ALL reports."""
        self.flow_exec.args = argparse.Namespace(report='ALL')
        result = list(self.flow_exec.generate_contexts())
        expected = [{'report_name': r} for r in config.reports.keys()]
        assert result == expected