pcb-afm-simulation-system/scripts/test_executor_cancel.py

#!/usr/bin/env python3
# -*- coding: ascii -*-
"""Tests for web-side task cancellation in the executor point loop.

Bug context (2026-09-04): the PlanDetail stop button never called the
backend cancel API, and the executor never polled task status between
points, so a "stopped" scan kept consuming Motor-CAD time. These tests
cover the fix:
  1. _is_task_cancelled returns False without requests (local mode).
  2. _is_task_cancelled returns False on network/backend errors.
  3. execute_task aborts between points when the web reports 'cancelled'
     and reports the partial results with status 'cancelled'.

Note: the module-level `requests` in task_executor may be None in this
dev environment (requests not installed), so tests patch the module
attribute itself rather than requests.get.

Run: python scripts/test_executor_cancel.py   (exit 0 = PASS)
"""
import os
import sys
import tempfile
import unittest
from unittest.mock import MagicMock, patch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import task_executor  # noqa: E402


class FakeResponse:
    def __init__(self, status_code=200, payload=None):
        self.status_code = status_code
        self._payload = payload or {}

    def json(self):
        return self._payload


def make_executor(**kwargs):
    kwargs.setdefault("web_base_url", "http://web.test")
    kwargs.setdefault("enable_mock", True)
    kwargs.setdefault("task_dir", tempfile.mkdtemp())
    return task_executor.TaskExecutor(**kwargs)


def make_task(n_points=5):
    return {
        "task_id": "t-cancel-1",
        "task_name": "cancel-test",
        "status": "dispatched",
        "parameters": [{"x": i} for i in range(n_points)],
    }


def mock_requests(get_side_effect=None):
    """Build a mock `requests` module with a controllable get()."""
    m = MagicMock()
    if get_side_effect is not None:
        m.get.side_effect = get_side_effect
    return m


class TestIsTaskCancelled(unittest.TestCase):
    def test_local_mode_never_cancelled(self):
        ex = make_executor()
        with patch.object(task_executor, "requests", None):
            self.assertFalse(ex._is_task_cancelled("t1"))

    def test_network_error_returns_false(self):
        ex = make_executor()
        m = mock_requests()
        m.get.side_effect = Exception("boom")
        with patch.object(task_executor, "requests", m):
            self.assertFalse(ex._is_task_cancelled("t1"))

    def test_cancelled_status_returns_true(self):
        ex = make_executor()
        m = mock_requests()
        m.get.return_value = FakeResponse(200, {"status": "cancelled"})
        with patch.object(task_executor, "requests", m):
            self.assertTrue(ex._is_task_cancelled("t1"))

    def test_running_status_returns_false(self):
        ex = make_executor()
        m = mock_requests()
        m.get.return_value = FakeResponse(200, {"status": "running"})
        with patch.object(task_executor, "requests", m):
            self.assertFalse(ex._is_task_cancelled("t1"))


class TestExecuteTaskCancellation(unittest.TestCase):
    def test_aborts_mid_run_and_reports_cancelled(self):
        """Cancel after the first point: only 1 point runs, status=cancelled."""
        ex = make_executor()
        task = make_task(n_points=5)
        run_params = []

        orig_run = ex._run_simulation_point

        def spy_run(params, index):
            run_params.append(index)
            return orig_run(params, index)

        ex._run_simulation_point = spy_run

        # GET sequence: (1) _hydrate_task downloads the task payload, (2)
        # status check before point 0 -> running, (3) check before point 1 ->
        # cancelled -> exactly one point executes.
        m = mock_requests()
        m.get.side_effect = [
            FakeResponse(200, {}),  # hydrate: no payload, keep local params
            FakeResponse(200, {"status": "running"}),
            FakeResponse(200, {"status": "cancelled"}),
        ]

        reported = {}

        def fake_report_results(task_id, results, metrics, logs, duration, status):
            reported["status"] = status
            reported["n_results"] = len(results)
            return True

        with patch.object(task_executor, "requests", m), \
             patch.object(ex, "report_results", side_effect=fake_report_results), \
             patch.object(ex, "report_progress", return_value=True):
            ex.execute_task(task)

        self.assertEqual(run_params, [0], "only the first point should run")
        self.assertEqual(reported["status"], "cancelled")
        self.assertEqual(reported["n_results"], 1)
        self.assertTrue(ex._task_cancelled)

    def test_not_cancelled_runs_all_points(self):
        ex = make_executor()
        task = make_task(n_points=3)
        reported = {}

        def fake_report_results(task_id, results, metrics, logs, duration, status):
            reported["status"] = status
            reported["n_results"] = len(results)
            return True

        m = mock_requests()
        m.get.return_value = FakeResponse(200, {"status": "running"})

        with patch.object(task_executor, "requests", m), \
             patch.object(ex, "report_results", side_effect=fake_report_results), \
             patch.object(ex, "report_progress", return_value=True):
            ex.execute_task(task)

        self.assertEqual(reported["status"], "completed")
        self.assertEqual(reported["n_results"], 3)
        self.assertFalse(ex._task_cancelled)


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