fix(product-suite): finalize and cancel generation

tests/test_ai.py

@@ -6,6 +6,7 @@ import socket
 import sys
 import threading
 import unittest
+from concurrent.futures import CancelledError
 from types import SimpleNamespace
 from unittest import mock
 
@@ -38,6 +39,7 @@ class _RequestsResponse:
         self.content = content
         self.headers = headers or {}
         self.text = json.dumps(self.payload, ensure_ascii=False)
+        self.closed = False
 
     def json(self):
         return self.payload
@@ -46,6 +48,9 @@ class _RequestsResponse:
         if self.content:
             yield self.content
 
+    def close(self):
+        self.closed = True
+
 class AITests(TempDirMixin, unittest.TestCase):
     def _write_models(self, path, text=None, image=None):
         text = text or {
@@ -732,7 +737,16 @@ class AITests(TempDirMixin, unittest.TestCase):
         ]
         calls = []
 
-        def fake_run(args, **kwargs):
+        class FakeProcess:
+            returncode = 0
+
+            def poll(self):
+                return self.returncode
+
+            def communicate(self):
+                return b"", b""
+
+        def fake_popen(args, **kwargs):
             calls.append((args, kwargs))
             self.assertIn("-K", args)
             config_path = args[args.index("-K") + 1]
@@ -755,12 +769,12 @@ class AITests(TempDirMixin, unittest.TestCase):
             output_path = args[args.index("--output") + 1]
             with open(output_path, "wb") as fh:
                 fh.write(generated_png)
-            return SimpleNamespace(returncode=0, stdout=b"", stderr=b"")
+            return FakeProcess()
 
         with mock.patch("app.ai.os.name", "nt"), \
                 mock.patch("app.ai.subprocess.CREATE_NO_WINDOW", 0x08000000, create=True), \
                 mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
-                mock.patch("app.ai.subprocess.run", side_effect=fake_run), \
+                mock.patch("app.ai.subprocess.Popen", side_effect=fake_popen), \
                 mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
             image_bytes = ai._download_cmhub_image(
                 url,
@@ -786,7 +800,7 @@ class AITests(TempDirMixin, unittest.TestCase):
 
     def test_cmhub_image_download_skips_curl_for_private_url(self):
         with mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
-                mock.patch("app.ai.subprocess.run") as run:
+                mock.patch("app.ai.subprocess.Popen") as popen:
             with self.assertRaises(ai.AIError):
                 ai._download_cmhub_image(
                     "http://127.0.0.1/a.png",
@@ -794,7 +808,7 @@ class AITests(TempDirMixin, unittest.TestCase):
                     read_timeout=ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
                     download_with_curl="true",
                 )
-        run.assert_not_called()
+        popen.assert_not_called()
 
     def test_cmhub_image_download_falls_back_to_requests_when_curl_fails(self):
         generated_png = self._png_bytes()
@@ -805,11 +819,20 @@ class AITests(TempDirMixin, unittest.TestCase):
         def fake_get(url, **kwargs):
             return _RequestsResponse(content=generated_png)
 
+        class FailedProcess:
+            returncode = 28
+
+            def poll(self):
+                return self.returncode
+
+            def communicate(self):
+                return b"", b"timeout"
+
         with mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
                 mock.patch(
-                    "app.ai.subprocess.run",
-                    return_value=SimpleNamespace(returncode=28, stdout=b"", stderr=b"timeout"),
-                ) as run, \
+                    "app.ai.subprocess.Popen",
+                    return_value=FailedProcess(),
+                ) as popen, \
                 mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get) as get, \
                 mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
             image_bytes = ai._download_cmhub_image(
@@ -820,7 +843,7 @@ class AITests(TempDirMixin, unittest.TestCase):
             )
 
         self.assertEqual(generated_png, image_bytes)
-        self.assertEqual(1, run.call_count)
+        self.assertEqual(1, popen.call_count)
         self.assertEqual(1, get.call_count)
 
     def test_cmhub_image_download_auto_without_curl_uses_requests(self):
@@ -834,7 +857,7 @@ class AITests(TempDirMixin, unittest.TestCase):
 
         with mock.patch("app.ai.os.name", "nt"), \
                 mock.patch("app.ai._find_system_curl", return_value=""), \
-                mock.patch("app.ai.subprocess.run") as run, \
+                mock.patch("app.ai.subprocess.Popen") as popen, \
                 mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
                 mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
             image_bytes = ai._download_cmhub_image(
@@ -845,7 +868,7 @@ class AITests(TempDirMixin, unittest.TestCase):
             )
 
         self.assertEqual(generated_png, image_bytes)
-        run.assert_not_called()
+        popen.assert_not_called()
 
     def test_cmhub_image_download_auto_on_non_windows_uses_requests(self):
         generated_png = self._png_bytes()
@@ -858,7 +881,7 @@ class AITests(TempDirMixin, unittest.TestCase):
 
         with mock.patch("app.ai.os.name", "posix"), \
                 mock.patch("app.ai._find_system_curl", return_value="/usr/bin/curl"), \
-                mock.patch("app.ai.subprocess.run") as run, \
+                mock.patch("app.ai.subprocess.Popen") as popen, \
                 mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
                 mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
             image_bytes = ai._download_cmhub_image(
@@ -869,7 +892,93 @@ class AITests(TempDirMixin, unittest.TestCase):
             )
 
         self.assertEqual(generated_png, image_bytes)
-        run.assert_not_called()
+        popen.assert_not_called()
+
+    def test_cmhub_requests_download_cancel_closes_response(self):
+        stopped = {"value": False}
+
+        class StreamingResponse(_RequestsResponse):
+            def iter_content(self, chunk_size=65536):
+                yield b"first"
+                stopped["value"] = True
+                yield b"second"
+
+        response = StreamingResponse()
+        with mock.patch.object(
+            ai._cmhub_session(),
+            "get",
+            return_value=response,
+        ):
+            with self.assertRaises(CancelledError):
+                ai._download_cmhub_image_with_requests(
+                    "https://cdn.example.com/generated.png",
+                    connect_timeout=3,
+                    read_timeout=30,
+                    max_bytes=ai.CMHUB_IMAGE_MAX_BYTES,
+                    should_stop=lambda: stopped["value"],
+                )
+
+        self.assertTrue(response.closed)
+
+    def test_cmhub_curl_download_cancel_terminates_process_and_cleans_temp_files(self):
+        stopped = {"value": False}
+        captured_paths = []
+
+        class RunningProcess:
+            def __init__(self):
+                self.returncode = None
+                self.terminated = False
+                self.killed = False
+
+            def poll(self):
+                return self.returncode
+
+            def terminate(self):
+                self.terminated = True
+                self.returncode = -15
+
+            def kill(self):
+                self.killed = True
+                self.returncode = -9
+
+            def wait(self, timeout=None):
+                return self.returncode
+
+            def communicate(self):
+                return b"", b""
+
+        process = RunningProcess()
+
+        def fake_popen(args, **kwargs):
+            captured_paths.extend(
+                [
+                    args[args.index("-K") + 1],
+                    args[args.index("--output") + 1],
+                ]
+            )
+            stopped["value"] = True
+            return process
+
+        with mock.patch(
+            "app.ai._find_system_curl",
+            return_value=r"C:\Windows\System32\curl.exe",
+        ), mock.patch(
+            "app.ai.subprocess.Popen",
+            side_effect=fake_popen,
+        ):
+            with self.assertRaises(CancelledError):
+                ai._download_cmhub_image_with_curl(
+                    "https://cdn.example.com/generated.png",
+                    connect_timeout=3,
+                    read_timeout=30,
+                    max_bytes=ai.CMHUB_IMAGE_MAX_BYTES,
+                    should_stop=lambda: stopped["value"],
+                )
+
+        self.assertTrue(process.terminated)
+        self.assertFalse(process.killed)
+        self.assertTrue(captured_paths)
+        self.assertTrue(all(not os.path.exists(path) for path in captured_paths))
 
     def test_cmhub_upstream_error_retries_and_keeps_metadata(self):
         with self.make_temp_dir() as temp_dir: