parsedmarc/tests/test_parallel.py

"""Tests for parsedmarc.parallel"""

import functools
import logging
import os
import tempfile
import unittest
from unittest.mock import patch

import parsedmarc
from parsedmarc.parallel import (
    _init_worker,
    _parse_report_email_job,
    _parse_report_file_job,
    parallel_map,
)

# Stable sample files reused from tests/test_cli.py's TestDirectoryFilePaths,
# plus a third plain aggregate sample, so parity/order tests exercise more
# than one worker submission.
SAMPLE_PATHS = [
    "samples/aggregate/!example.com!1538204542!1538463818.xml",
    "samples/aggregate/!large-example.com!1711897200!1711983600.xml",
    "samples/aggregate/example.net!example.com!1529366400!1529452799.xml",
]


def _echo_job(x):
    """Trivial module-level (spawn-picklable) worker used to exercise
    parallel_map's scheduling behavior without involving real parsing."""
    return x


class _CountingIterable:
    """Wraps a range so tests can observe how many items parallel_map has
    pulled from a lazily-iterated jobs source at any point during
    iteration, without materializing the whole sequence up front."""

    def __init__(self, n):
        self.n = n
        self.pulled = 0

    def __iter__(self):
        for i in range(self.n):
            self.pulled += 1
            yield i


class _ParallelTestCase(unittest.TestCase):
    """Common env setup shared by parallel.py tests: offline mode, no DNS,
    and a cold IP address cache."""

    def setUp(self):
        self._env_patcher = patch.dict(
            os.environ, {"GITHUB_ACTIONS": "true"}, clear=False
        )
        self._env_patcher.start()
        self.addCleanup(self._env_patcher.stop)
        # Earlier test modules (e.g. tests/test_init.py) may run with DNS
        # enabled locally and warm the shared module-level
        # parsedmarc.IP_ADDRESS_CACHE. get_ip_address_info consults the
        # cache before the offline check, so a warm cache would leak
        # DNS-enriched entries into offline parses run in this process
        # while spawned workers start cold, breaking parity.
        parsedmarc.IP_ADDRESS_CACHE.clear()


class TestParallelMapParseReportFile(_ParallelTestCase):
    """parallel_map + _parse_report_file_job over real sample files must
    behave identically to calling parse_report_file sequentially, and
    must preserve submission order in its results."""

    def test_results_match_sequential_parsing_in_order(self):
        expected = [
            (path, parsedmarc.parse_report_file(path, offline=True))
            for path in SAMPLE_PATHS
        ]

        job = functools.partial(
            _parse_report_file_job, config=parsedmarc.ParserConfig(offline=True)
        )
        results = list(parallel_map(job, SAMPLE_PATHS, n_procs=2))

        self.assertEqual(len(results), len(SAMPLE_PATHS))
        # Order must match submission order (SAMPLE_PATHS), not completion
        # order.
        self.assertEqual([path for path, _ in results], SAMPLE_PATHS)
        for (path, report), (expected_path, expected_report) in zip(results, expected):
            self.assertEqual(path, expected_path)
            self.assertNotIsInstance(report, Exception)
            self.assertEqual(report, expected_report)


class TestParallelMapJunkFile(_ParallelTestCase):
    """A worker crash on an unparseable file must surface as an Exception
    *value* in the result tuple, not hang the whole run. This is a
    regression guard for the old Pipe/Process CLI worker, which left the
    parent blocked forever on conn.recv() when a child died from a
    non-ParserError exception."""

    def test_junk_file_yields_exception_value_and_completes(self):
        with tempfile.NamedTemporaryFile(suffix=".xml", delete=False, mode="wb") as tf:
            tf.write(b"not a report")
            junk_path = tf.name
        self.addCleanup(os.remove, junk_path)

        job = functools.partial(
            _parse_report_file_job, config=parsedmarc.ParserConfig(offline=True)
        )
        results = list(parallel_map(job, [junk_path, junk_path], n_procs=2))

        self.assertEqual(len(results), 2)
        for path, result in results:
            self.assertEqual(path, junk_path)
            self.assertIsInstance(result, Exception)


class TestParallelMapBoundedLaziness(unittest.TestCase):
    """The jobs iterable must never be materialized up front. At any point
    during iteration, the number of items pulled from the source should
    stay within window_factor * n_procs of the number of results already
    yielded, bounding memory use for very large inputs (e.g. a 20,000
    message mbox)."""

    def test_consumption_stays_bounded_and_results_are_complete_and_ordered(self):
        n = 20
        window_factor = 2
        n_procs = 2
        window_size = window_factor * n_procs

        jobs = _CountingIterable(n)
        results = []
        gen = parallel_map(
            _echo_job, jobs, n_procs=n_procs, window_factor=window_factor
        )
        for result in gen:
            results.append(result)
            # +1 buffer: the generator may pull one extra job before it
            # can submit-then-harvest on a given step.
            self.assertLessEqual(jobs.pulled, window_size + len(results) + 1)

        self.assertEqual(results, list(range(n)))


class TestParallelMapShouldStop(unittest.TestCase):
    """should_stop lets a caller (e.g. a CLI handling SIGTERM) end a run
    early without raising, while still yielding results already
    completed in the submission window."""

    def test_should_stop_ends_iteration_early(self):
        n = 20
        jobs = _CountingIterable(n)

        results = list(
            parallel_map(_echo_job, jobs, n_procs=2, should_stop=lambda: True)
        )

        self.assertLess(len(results), n)
        self.assertLess(jobs.pulled, n)
        # Results still seen so far must be a prefix of submission order.
        self.assertEqual(results, list(range(len(results))))


class TestParallelMapValidation(unittest.TestCase):
    """parallel_map is a reusable helper, so it validates n_procs itself
    with a clear message instead of surfacing ProcessPoolExecutor's
    max_workers error later - and it must do so eagerly at the call, not
    on first iteration of the returned iterator (callers that pass the
    iterator elsewhere before consuming it would otherwise see the error
    far from the bad argument)."""

    def test_n_procs_below_one_raises_value_error_eagerly(self):
        with self.assertRaises(ValueError):
            parallel_map(_echo_job, [1, 2], n_procs=0)


class TestParallelMapEmptyJobs(unittest.TestCase):
    """An empty jobs iterable must return immediately without spawning a
    process pool."""

    def test_empty_jobs_yields_nothing_and_spawns_no_pool(self):
        with patch("parsedmarc.parallel.ProcessPoolExecutor") as mock_executor:
            results = list(parallel_map(_echo_job, [], n_procs=2))
        self.assertEqual(results, [])
        mock_executor.assert_not_called()


class TestWorkerLogging(_ParallelTestCase):
    """Worker processes must reconstruct the parent parsedmarc logger's
    level and FileHandler(s) so records emitted during parsing (e.g.
    parse_report_file's "Parsing <path>" debug line) aren't silently
    dropped just because they happened in a child process."""

    def setUp(self):
        super().setUp()
        from parsedmarc.log import logger as plog

        self._saved_handlers = list(plog.handlers)
        self._saved_level = plog.level

    def tearDown(self):
        from parsedmarc.log import logger as plog

        for handler in list(plog.handlers):
            if handler not in self._saved_handlers:
                plog.removeHandler(handler)
                if isinstance(handler, logging.FileHandler):
                    handler.close()
        plog.handlers[:] = self._saved_handlers
        plog.setLevel(self._saved_level)
        super().tearDown()

    def test_worker_debug_log_reaches_parent_file_handler(self):
        from parsedmarc.log import configure_logging

        with tempfile.NamedTemporaryFile(suffix=".log", delete=False) as tf:
            log_path = tf.name
        self.addCleanup(lambda: os.path.exists(log_path) and os.remove(log_path))

        configure_logging(logging.DEBUG, log_path)

        sample = SAMPLE_PATHS[0]
        job = functools.partial(
            _parse_report_file_job, config=parsedmarc.ParserConfig(offline=True)
        )
        results = list(parallel_map(job, [sample], n_procs=2))

        self.assertEqual(len(results), 1)
        path, report = results[0]
        self.assertEqual(path, sample)
        self.assertNotIsInstance(report, Exception)

        with open(log_path) as f:
            contents = f.read()
        # parse_report_file logs `Parsing {file_path}` at DEBUG
        # (parsedmarc/__init__.py) -- this line only appears if the
        # worker process's reconstructed logger actually wrote to the
        # parent's log file.
        self.assertIn(f"Parsing {sample}", contents)


class TestInitWorkerLogging(unittest.TestCase):
    """_init_worker must be usable directly as a pool initializer: with no
    log files it sets the level (configure_logging still ensures a console
    handler), and with log files it attaches a FileHandler per path. The
    IP-lookup state and PSL overrides it also restores are passed through
    unchanged here so these tests exercise only the logging side."""

    def setUp(self):
        from parsedmarc.log import logger as plog

        self._saved_handlers = list(plog.handlers)
        self._saved_level = plog.level

    def tearDown(self):
        from parsedmarc.log import logger as plog

        for handler in list(plog.handlers):
            if handler not in self._saved_handlers:
                plog.removeHandler(handler)
                if isinstance(handler, logging.FileHandler):
                    handler.close()
        plog.handlers[:] = self._saved_handlers
        plog.setLevel(self._saved_level)

    def test_no_log_files_sets_level_only(self):
        from parsedmarc import utils
        from parsedmarc.log import logger as plog

        _init_worker(
            logging.WARNING,
            [],
            utils._get_ip_worker_state(),
            list(utils.psl_overrides),
        )
        self.assertEqual(plog.level, logging.WARNING)

    def test_log_files_attach_file_handlers(self):
        from parsedmarc import utils
        from parsedmarc.log import logger as plog

        with tempfile.NamedTemporaryFile(suffix=".log", delete=False) as tf:
            log_path = tf.name
        self.addCleanup(lambda: os.path.exists(log_path) and os.remove(log_path))

        _init_worker(
            logging.DEBUG,
            [log_path],
            utils._get_ip_worker_state(),
            list(utils.psl_overrides),
        )

        self.assertEqual(plog.level, logging.DEBUG)
        file_handlers = [h for h in plog.handlers if isinstance(h, logging.FileHandler)]
        self.assertTrue(any(h.baseFilename == log_path for h in file_handlers))
        for h in file_handlers:
            h.close()


def _ip_worker_state_job(_):
    """Read the worker process's IP-lookup state and PSL overrides so a
    test can compare them with what the parent configured. Module-level and
    picklable (a spawn child re-imports this module to unpickle it)."""
    from parsedmarc import utils

    return (utils._IPINFO_API_TOKEN, utils._IP_DB_PATH, list(utils.psl_overrides))


class TestWorkerIPState(unittest.TestCase):
    """parallel_map workers must make the same IPinfo API / MMDB selections
    and use the same PSL overrides the parent configured, whatever the
    start method.

    parsedmarc.utils._IPINFO_API_TOKEN (configure_ipinfo_api),
    parsedmarc.utils._IP_DB_PATH (load_ip_db), and the module-level
    parsedmarc.utils.psl_overrides list (load_psl_overrides) are process
    globals. Under ``fork`` a worker inherits them ("All resources of the
    parent are inherited by the child process"), but under ``spawn`` and
    ``forkserver`` it does not ("The child process will only inherit those
    resources necessary to run the process object's ``run()`` method") -
    https://docs.python.org/3/library/multiprocessing.html. The pool
    initializer therefore restores the parent's snapshot in every worker.
    This test forces each non-fork start method so it guards that path on
    every Python version, not only the ones whose default start method
    already exercises it (Linux defaults to forkserver from 3.14, macOS
    and Windows to spawn).

    With the ``psl_overrides`` restore removed from ``_init_worker``, the
    PSL assertion failed under both ``spawn`` and ``forkserver``: each
    worker reported the bundled ``psl_overrides.txt`` list it loads at
    import, not the parent's ``["example.test"]``.
    """

    def test_workers_see_parent_ip_state_under_spawn_and_forkserver(self):
        import multiprocessing

        import parsedmarc.utils as utils

        methods = [
            m
            for m in ("spawn", "forkserver")
            if m in multiprocessing.get_all_start_methods()
        ]
        self.assertIn("spawn", methods)

        original_state = utils._get_ip_worker_state()
        original_psl = list(utils.psl_overrides)
        original_method = multiprocessing.get_start_method()
        utils._set_ip_worker_state(
            "worker-state-test-token", "/parent/only/ipinfo_lite.mmdb"
        )
        utils.psl_overrides[:] = ["example.test"]
        self.addCleanup(utils._set_ip_worker_state, *original_state)
        self.addCleanup(utils.psl_overrides.__setitem__, slice(None), original_psl)
        self.addCleanup(multiprocessing.set_start_method, original_method, force=True)

        for method in methods:
            with self.subTest(start_method=method):
                multiprocessing.set_start_method(method, force=True)

                results = list(parallel_map(_ip_worker_state_job, range(8), n_procs=2))

                self.assertEqual(len(results), 8)
                for token, db_path, psl in results:
                    self.assertEqual(token, "worker-state-test-token")
                    self.assertEqual(db_path, "/parent/only/ipinfo_lite.mmdb")
                    self.assertEqual(psl, ["example.test"])


class TestParseReportEmailJob(_ParallelTestCase):
    """_parse_report_email_job must return a ParserError as a value (never
    raise it) on an invalid message."""

    def test_invalid_email_returns_parser_error_value(self):
        result = _parse_report_email_job(
            b"not a valid email", config=parsedmarc.ParserConfig(offline=True)
        )
        self.assertIsInstance(result, parsedmarc.ParserError)


if __name__ == "__main__":
    unittest.main(verbosity=2)