import asyncio
import io
import json
import os
import sys
import tempfile
import unittest
from types import ModuleType, SimpleNamespace
from unittest.mock import MagicMock, patch

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from python import persistent_wrapper


class TestCallMethod(unittest.TestCase):
    """Tests for call_method function."""

    def test_sync_function(self) -> None:
        """Tests call_method with a synchronous function."""

        def sync_func(a, b):
            return a + b

        result = persistent_wrapper.call_method(sync_func, [2, 3])
        self.assertEqual(result, 5)

    def test_async_function(self) -> None:
        """Tests call_method with an asynchronous function."""

        async def async_func(a, b):
            await asyncio.sleep(0.001)
            return a * b

        result = persistent_wrapper.call_method(async_func, [3, 4])
        self.assertEqual(result, 12)

    def test_detects_async_correctly(self) -> None:
        """Tests that call_method uses inspect.iscoroutinefunction (not deprecated asyncio version)."""

        async def real_async_func():
            return "async"

        def real_sync_func():
            return "sync"

        # Test async function is detected and run with asyncio.run
        def close_coroutine(coro):
            coro.close()
            return "mocked_async"

        with patch("asyncio.run", side_effect=close_coroutine) as mock_run:
            result = persistent_wrapper.call_method(real_async_func, [])
            mock_run.assert_called_once()
            self.assertEqual(result, "mocked_async")

        # Test sync function bypasses asyncio.run
        with patch("asyncio.run") as mock_run:
            result = persistent_wrapper.call_method(real_sync_func, [])
            mock_run.assert_not_called()
            self.assertEqual(result, "sync")


class TestLoadUserModule(unittest.TestCase):
    """Tests for load_user_module function."""

    def test_loads_valid_module(self) -> None:
        """Tests loading a user module from a file."""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write("def test_func(x):\n    return x * 2\n")
            f.flush()
            script_path = f.name

        try:
            module = persistent_wrapper.load_user_module(script_path)
            self.assertTrue(hasattr(module, "test_func"))
            self.assertEqual(module.test_func(5), 10)
        finally:
            os.unlink(script_path)

    def test_file_not_found(self) -> None:
        """Tests loading a non-existent module raises FileNotFoundError."""
        with self.assertRaises(FileNotFoundError):
            persistent_wrapper.load_user_module("/nonexistent/path.py")

    def test_invalid_spec_raises_import_error(self) -> None:
        """Tests that ImportError is raised when spec is None."""
        with patch("importlib.util.spec_from_file_location", return_value=None):
            with self.assertRaises(ImportError) as context:
                persistent_wrapper.load_user_module("/some/path.py")
            self.assertIn("Cannot load module", str(context.exception))

    def test_adds_script_dir_to_path(self) -> None:
        """Tests that script directory is added to sys.path."""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write("x = 1\n")
            f.flush()
            script_path = f.name
            script_dir = os.path.dirname(os.path.abspath(script_path))

        try:
            # Remove from path if present
            if script_dir in sys.path:
                sys.path.remove(script_dir)

            persistent_wrapper.load_user_module(script_path)
            self.assertIn(script_dir, sys.path)
        finally:
            os.unlink(script_path)
            if script_dir in sys.path:
                sys.path.remove(script_dir)


class TestGetCallable(unittest.TestCase):
    """Tests for get_callable function."""

    def test_gets_top_level_function(self) -> None:
        """Tests getting a top-level function from a module."""
        mock_module = MagicMock()
        mock_module.my_func = lambda x: x

        result = persistent_wrapper.get_callable(mock_module, "my_func")
        self.assertEqual(result, mock_module.my_func)

    def test_gets_class_method(self) -> None:
        """Tests getting a class method using 'Class.method' syntax."""
        mock_module = MagicMock()
        mock_module.MyClass.my_method = lambda x: x

        result = persistent_wrapper.get_callable(mock_module, "MyClass.my_method")
        self.assertEqual(result, mock_module.MyClass.my_method)

    def test_not_found_shows_available_functions(self) -> None:
        """Tests that AttributeError shows available functions."""
        mock_module = MagicMock(spec=["existing_func"])
        mock_module.__name__ = "test_module"
        mock_module.existing_func = lambda: None

        with self.assertRaises(AttributeError) as context:
            persistent_wrapper.get_callable(mock_module, "nonexistent_func")

        error_message = str(context.exception)
        self.assertIn("nonexistent_func", error_message)
        self.assertIn("not found", error_message)

    def test_suggests_get_to_call_rename(self) -> None:
        """Tests suggestion when user has get_api instead of call_api."""
        # Create a real module-like object
        mock_module = type("MockModule", (), {"__name__": "test_module"})()
        mock_module.get_api = lambda: None

        with self.assertRaises(AttributeError) as context:
            persistent_wrapper.get_callable(mock_module, "call_api")

        error_message = str(context.exception)
        self.assertIn("Did you mean to rename", error_message)
        self.assertIn("get_api", error_message)

    def test_suggests_missing_call_prefix(self) -> None:
        """Tests suggestion when user has 'api' instead of 'call_api'."""
        mock_module = type("MockModule", (), {"__name__": "test_module"})()
        mock_module.api = lambda: None

        with self.assertRaises(AttributeError) as context:
            persistent_wrapper.get_callable(mock_module, "call_api")

        error_message = str(context.exception)
        self.assertIn("Did you mean to rename", error_message)
        self.assertIn("api", error_message)

    def test_suggests_typo_correction(self) -> None:
        """Tests suggestion for typo-like mistakes."""
        mock_module = type("MockModule", (), {"__name__": "test_module"})()
        mock_module.call_apii = lambda: None  # typo: extra 'i'

        with self.assertRaises(AttributeError) as context:
            persistent_wrapper.get_callable(mock_module, "call_api")

        error_message = str(context.exception)
        self.assertIn("Did you mean", error_message)
        self.assertIn("call_apii", error_message)

    def test_no_suggestion_when_no_similar_function(self) -> None:
        """Tests no suggestion when no similar function exists."""
        mock_module = type("MockModule", (), {"__name__": "test_module"})()
        mock_module.completely_different = lambda: None

        with self.assertRaises(AttributeError) as context:
            persistent_wrapper.get_callable(mock_module, "call_api")

        error_message = str(context.exception)
        self.assertIn("call_api", error_message)
        self.assertNotIn("Did you mean", error_message)


class TestInitTracing(unittest.TestCase):
    """Tests for OpenTelemetry initialization."""

    def setUp(self) -> None:
        self.original_tracer = persistent_wrapper._tracer
        self.original_tracing_enabled = persistent_wrapper._tracing_enabled
        persistent_wrapper._tracer = None
        persistent_wrapper._tracing_enabled = False

    def tearDown(self) -> None:
        persistent_wrapper._tracer = self.original_tracer
        persistent_wrapper._tracing_enabled = self.original_tracing_enabled

    def test_successful_init_does_not_write_to_stderr(self) -> None:
        """Tests that successful tracing initialization stays quiet."""
        fake_trace = MagicMock()
        fake_tracer = object()
        fake_trace.get_tracer.return_value = fake_tracer

        opentelemetry_module = ModuleType("opentelemetry")
        opentelemetry_exporter_module = ModuleType("opentelemetry.exporter")
        opentelemetry_exporter_otlp_module = ModuleType("opentelemetry.exporter.otlp")
        opentelemetry_exporter_otlp_proto_module = ModuleType(
            "opentelemetry.exporter.otlp.proto"
        )
        opentelemetry_exporter_otlp_proto_http_module = ModuleType(
            "opentelemetry.exporter.otlp.proto.http"
        )
        trace_exporter_module = ModuleType(
            "opentelemetry.exporter.otlp.proto.http.trace_exporter"
        )
        sdk_module = ModuleType("opentelemetry.sdk")
        sdk_trace_module = ModuleType("opentelemetry.sdk.trace")
        sdk_trace_export_module = ModuleType("opentelemetry.sdk.trace.export")

        class FakeOTLPSpanExporter:
            def __init__(self, endpoint):
                self.endpoint = endpoint

        class FakeSimpleSpanProcessor:
            def __init__(self, exporter):
                self.exporter = exporter

        class FakeTracerProvider:
            def __init__(self):
                self.processors = []

            def add_span_processor(self, processor):
                self.processors.append(processor)

        trace_exporter_module.OTLPSpanExporter = FakeOTLPSpanExporter
        sdk_trace_module.TracerProvider = FakeTracerProvider
        sdk_trace_export_module.SimpleSpanProcessor = FakeSimpleSpanProcessor

        opentelemetry_module.trace = fake_trace
        opentelemetry_module.exporter = opentelemetry_exporter_module
        opentelemetry_module.sdk = sdk_module
        opentelemetry_exporter_module.otlp = opentelemetry_exporter_otlp_module
        opentelemetry_exporter_otlp_module.proto = (
            opentelemetry_exporter_otlp_proto_module
        )
        opentelemetry_exporter_otlp_proto_module.http = (
            opentelemetry_exporter_otlp_proto_http_module
        )
        opentelemetry_exporter_otlp_proto_http_module.trace_exporter = (
            trace_exporter_module
        )
        sdk_module.trace = sdk_trace_module
        sdk_trace_module.export = sdk_trace_export_module

        def getenv(name, default=None):
            if name == "PROMPTFOO_ENABLE_OTEL":
                return "true"
            return default

        fake_modules = {
            "opentelemetry": opentelemetry_module,
            "opentelemetry.exporter": opentelemetry_exporter_module,
            "opentelemetry.exporter.otlp": opentelemetry_exporter_otlp_module,
            "opentelemetry.exporter.otlp.proto": opentelemetry_exporter_otlp_proto_module,
            "opentelemetry.exporter.otlp.proto.http": (
                opentelemetry_exporter_otlp_proto_http_module
            ),
            "opentelemetry.exporter.otlp.proto.http.trace_exporter": (
                trace_exporter_module
            ),
            "opentelemetry.sdk": sdk_module,
            "opentelemetry.sdk.trace": sdk_trace_module,
            "opentelemetry.sdk.trace.export": sdk_trace_export_module,
        }

        with patch.dict(sys.modules, fake_modules, clear=False):
            with patch("python.persistent_wrapper.os.getenv", side_effect=getenv):
                with patch("sys.stderr", new_callable=io.StringIO) as stderr_capture:
                    persistent_wrapper._init_tracing()

        self.assertTrue(persistent_wrapper._tracing_enabled)
        self.assertIs(persistent_wrapper._tracer, fake_tracer)
        self.assertEqual(stderr_capture.getvalue(), "")


class TestTracedCall(unittest.TestCase):
    """Tests standardized OpenTelemetry attributes on traced Python provider calls."""

    def test_uses_current_provider_and_usage_attribute_names(self) -> None:
        span = MagicMock()
        tracer = MagicMock()
        tracer.start_as_current_span.return_value.__enter__.return_value = span

        propagate_module = ModuleType("opentelemetry.propagate")
        propagate_module.extract = MagicMock(return_value=object())
        trace_module = ModuleType("opentelemetry.trace")
        trace_module.SpanKind = SimpleNamespace(CLIENT="client")
        trace_module.Status = MagicMock()
        trace_module.StatusCode = SimpleNamespace(OK="ok", ERROR="error")

        with patch.dict(
            sys.modules,
            {
                "opentelemetry.propagate": propagate_module,
                "opentelemetry.trace": trace_module,
            },
        ):
            with patch.object(persistent_wrapper, "_tracer", tracer):
                with patch.object(persistent_wrapper, "_tracing_enabled", True):
                    result = persistent_wrapper._traced_call(
                        lambda *_args: {
                            "output": "done",
                            "tokenUsage": {
                                "prompt": 10,
                                "completion": 5,
                                "total": 15,
                            },
                        },
                        [
                            "prompt",
                            {"config": {"model": "customer-model"}},
                            {
                                "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01"
                            },
                        ],
                        "call_api",
                    )

        self.assertEqual(result["output"], "done")
        span.set_attribute.assert_any_call("promptfoo.provider.type", "python")
        span.set_attribute.assert_any_call("promptfoo.provider.function", "call_api")
        span.set_attribute.assert_any_call("promptfoo.provider.model", "customer-model")
        span.set_attribute.assert_any_call("gen_ai.usage.input_tokens", 10)
        span.set_attribute.assert_any_call("gen_ai.usage.output_tokens", 5)
        span.set_attribute.assert_any_call("promptfoo.usage.total_tokens", 15)
        attribute_names = [call.args[0] for call in span.set_attribute.call_args_list]
        self.assertNotIn("gen_ai.system", attribute_names)
        self.assertNotIn("gen_ai.operation.name", attribute_names)
        self.assertNotIn("gen_ai.provider.name", attribute_names)
        self.assertNotIn("gen_ai.request.model", attribute_names)
        self.assertNotIn("gen_ai.usage.total_tokens", attribute_names)


class TestMain(unittest.TestCase):
    """Tests for main function."""

    def test_exits_with_too_few_arguments(self) -> None:
        """Tests that main exits with error when not enough arguments."""
        with patch.object(sys, "argv", ["persistent_wrapper.py"]):
            with patch("sys.exit", side_effect=SystemExit(1)) as mock_exit:
                with patch("sys.stderr", new_callable=io.StringIO):
                    with self.assertRaises(SystemExit):
                        persistent_wrapper.main()
                    mock_exit.assert_called_with(1)

    def test_exits_on_module_load_failure(self) -> None:
        """Tests that main exits when module fails to load."""
        with patch.object(
            sys, "argv", ["persistent_wrapper.py", "/nonexistent.py", "func"]
        ):
            with patch("sys.exit", side_effect=SystemExit(1)) as mock_exit:
                with patch("sys.stderr", new_callable=io.StringIO):
                    with self.assertRaises(SystemExit):
                        persistent_wrapper.main()
                    mock_exit.assert_called_with(1)

    def test_signals_ready_and_handles_shutdown(self) -> None:
        """Tests READY signal and SHUTDOWN command."""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write("def test_func(): pass\n")
            f.flush()
            script_path = f.name

        try:
            stdout_capture = io.StringIO()
            stdin_mock = io.StringIO("SHUTDOWN\n")

            with patch.object(
                sys, "argv", ["persistent_wrapper.py", script_path, "test_func"]
            ):
                with patch("sys.stdin", stdin_mock):
                    with patch("sys.stdout", stdout_capture):
                        persistent_wrapper.main()

            output = stdout_capture.getvalue()
            self.assertIn("READY", output)
        finally:
            os.unlink(script_path)

    def test_ready_signal_survives_partial_stdout_during_import(self) -> None:
        """Regression (#10097): a user module that leaves stdout mid-line at
        import time must not clobber the READY marker - Node matches it as a
        whole line."""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write('import sys\nsys.stdout.write("loading")\ndef test_func(): pass\n')
            f.flush()
            script_path = f.name

        try:
            stdout_capture = io.StringIO()
            stdin_mock = io.StringIO("SHUTDOWN\n")

            with patch.object(
                sys, "argv", ["persistent_wrapper.py", script_path, "test_func"]
            ):
                with patch("sys.stdin", stdin_mock):
                    with patch("sys.stdout", stdout_capture):
                        persistent_wrapper.main()

            self.assertIn("READY", stdout_capture.getvalue().splitlines())
        finally:
            os.unlink(script_path)

    def test_handles_stdin_close(self) -> None:
        """Tests graceful exit when stdin is closed."""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write("def test_func(): pass\n")
            f.flush()
            script_path = f.name

        try:
            stdout_capture = io.StringIO()
            stdin_mock = io.StringIO("")  # Empty = EOF

            with patch.object(
                sys, "argv", ["persistent_wrapper.py", script_path, "test_func"]
            ):
                with patch("sys.stdin", stdin_mock):
                    with patch("sys.stdout", stdout_capture):
                        persistent_wrapper.main()

            output = stdout_capture.getvalue()
            self.assertIn("READY", output)
        finally:
            os.unlink(script_path)

    def test_handles_unknown_command(self) -> None:
        """Tests error message for unknown commands."""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write("def test_func(): pass\n")
            f.flush()
            script_path = f.name

        try:
            stderr_capture = io.StringIO()
            stdin_mock = io.StringIO("UNKNOWN_CMD\nSHUTDOWN\n")

            with patch.object(
                sys, "argv", ["persistent_wrapper.py", script_path, "test_func"]
            ):
                with patch("sys.stdin", stdin_mock):
                    with patch("sys.stderr", stderr_capture):
                        with patch("sys.stdout", io.StringIO()):
                            persistent_wrapper.main()

            stderr_output = stderr_capture.getvalue()
            self.assertIn("Unknown command", stderr_output)
        finally:
            os.unlink(script_path)

    def test_routes_call_command(self) -> None:
        """Tests that CALL commands are routed to handle_call."""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write("def test_func(): return 'ok'\n")
            f.flush()
            script_path = f.name

        try:
            with patch("python.persistent_wrapper.handle_call") as mock_handle_call:
                stdin_mock = io.StringIO("CALL|func|req|resp\nSHUTDOWN\n")

                with patch.object(
                    sys, "argv", ["persistent_wrapper.py", script_path, "test_func"]
                ):
                    with patch("sys.stdin", stdin_mock):
                        with patch("sys.stdout", io.StringIO()):
                            persistent_wrapper.main()

                mock_handle_call.assert_called_once()
        finally:
            os.unlink(script_path)


class TestHandleCall(unittest.TestCase):
    """Tests for handle_call function."""

    def setUp(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        self.request_file = os.path.join(self.temp_dir, "request.json")
        self.response_file = os.path.join(self.temp_dir, "response.json")

        # Create a mock module
        self.mock_module = type("MockModule", (), {"__name__": "test_module"})()
        self.mock_module.test_func = lambda a, b: a + b
        self.mock_module.error_func = self._raise_error

    def _raise_error(self):
        raise ValueError("Test error")

    def tearDown(self):
        """Clean up temp files."""
        import shutil

        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_new_format_success(self) -> None:
        """Tests successful CALL with new 4-part format."""
        with open(self.request_file, "w") as f:
            json.dump([2, 3], f)

        stdout_capture = io.StringIO()
        with patch("sys.stdout", stdout_capture):
            persistent_wrapper.handle_call(
                f"CALL|test_func|{self.request_file}|{self.response_file}",
                self.mock_module,
                "default_func",
            )

        self.assertIn("DONE", stdout_capture.getvalue())

        with open(self.response_file) as f:
            response = json.load(f)
        self.assertEqual(response["type"], "result")
        self.assertEqual(response["data"], 5)

    def test_done_marker_starts_own_line_after_partial_stdout(self) -> None:
        """Regression (#10097): user code that leaves stdout mid-line must not
        clobber the DONE| marker - Node matches it line-anchored."""

        def partial_stdout_func():
            sys.stdout.write("partial provider output")
            return {"ok": True}

        self.mock_module.partial_stdout_func = partial_stdout_func
        with open(self.request_file, "w") as f:
            json.dump([], f)

        stdout_capture = io.StringIO()
        with patch("sys.stdout", stdout_capture):
            persistent_wrapper.handle_call(
                f"CALL|partial_stdout_func|{self.request_file}|{self.response_file}",
                self.mock_module,
                "default_func",
            )

        done_lines = [
            line
            for line in stdout_capture.getvalue().splitlines()
            if line.startswith("DONE|")
        ]
        self.assertEqual(done_lines, [f"DONE|{self.response_file}"])

    def test_legacy_format_success(self) -> None:
        """Tests successful CALL with legacy 3-part format."""
        with open(self.request_file, "w") as f:
            json.dump([10, 20], f)

        stdout_capture = io.StringIO()
        with patch("sys.stdout", stdout_capture):
            persistent_wrapper.handle_call(
                f"CALL|{self.request_file}|{self.response_file}",
                self.mock_module,
                "test_func",  # Uses default function
            )

        self.assertIn("DONE", stdout_capture.getvalue())

        with open(self.response_file) as f:
            response = json.load(f)
        self.assertEqual(response["type"], "result")
        self.assertEqual(response["data"], 30)

    def test_invalid_format_writes_error(self) -> None:
        """Tests that invalid command format reports the error without emitting a bogus DONE marker."""
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()

        with patch("sys.stdout", stdout_capture):
            with patch("sys.stderr", stderr_capture):
                persistent_wrapper.handle_call(
                    "CALL|only_two_parts",
                    self.mock_module,
                    "test_func",
                )

        # No response_file means no useful DONE| can be emitted; the stderr is
        # the real signal and Node will resolve the call via timeout.
        self.assertNotIn("DONE", stdout_capture.getvalue())
        self.assertIn("Invalid CALL command format", stderr_capture.getvalue())

    def test_function_execution_error(self) -> None:
        """Tests that function errors are captured in response."""
        with open(self.request_file, "w") as f:
            json.dump([], f)

        stdout_capture = io.StringIO()
        with patch("sys.stdout", stdout_capture):
            persistent_wrapper.handle_call(
                f"CALL|error_func|{self.request_file}|{self.response_file}",
                self.mock_module,
                "default_func",
            )

        self.assertIn("DONE", stdout_capture.getvalue())

        with open(self.response_file) as f:
            response = json.load(f)
        self.assertEqual(response["type"], "error")
        self.assertIn("Test error", response["error"])

    def test_nonexistent_function_writes_error(self) -> None:
        """Tests that missing function writes error response."""
        with open(self.request_file, "w") as f:
            json.dump([], f)

        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()

        with patch("sys.stdout", stdout_capture):
            with patch("sys.stderr", stderr_capture):
                persistent_wrapper.handle_call(
                    f"CALL|nonexistent|{self.request_file}|{self.response_file}",
                    self.mock_module,
                    "default_func",
                )

        self.assertIn("DONE", stdout_capture.getvalue())

        with open(self.response_file) as f:
            response = json.load(f)
        self.assertEqual(response["type"], "error")

    def test_handles_unicode_in_response(self) -> None:
        """Tests that Unicode characters are preserved in response."""
        self.mock_module.unicode_func = lambda: {"emoji": "🎉", "chinese": "中文"}

        with open(self.request_file, "w") as f:
            json.dump([], f)

        with patch("sys.stdout", io.StringIO()):
            persistent_wrapper.handle_call(
                f"CALL|unicode_func|{self.request_file}|{self.response_file}",
                self.mock_module,
                "default_func",
            )

        with open(self.response_file, encoding="utf-8") as f:
            response = json.load(f)

        self.assertEqual(response["data"]["emoji"], "🎉")
        self.assertEqual(response["data"]["chinese"], "中文")

    def test_file_verification_retry(self) -> None:
        """Tests that file verification retries on failure."""
        with open(self.request_file, "w") as f:
            json.dump([1, 2], f)

        # Mock open to fail first two reads, succeed on third
        original_open = open
        verification_read_count = 0

        def mock_open_func(path, *args, **kwargs):
            nonlocal verification_read_count
            mode = args[0] if args else kwargs.get("mode", "r")
            if path == self.response_file and "r" in mode:
                verification_read_count += 1
                if verification_read_count <= 2:
                    raise IOError("Simulated read failure")
            return original_open(path, *args, **kwargs)

        stdout_capture = io.StringIO()
        with patch("sys.stdout", stdout_capture):
            with patch("time.sleep") as mock_sleep:
                with patch("builtins.open", side_effect=mock_open_func):
                    persistent_wrapper.handle_call(
                        f"CALL|test_func|{self.request_file}|{self.response_file}",
                        self.mock_module,
                        "default_func",
                    )

        self.assertIn("DONE", stdout_capture.getvalue())
        self.assertEqual(verification_read_count, 3)
        self.assertEqual(mock_sleep.call_count, 2)

        with open(self.response_file) as f:
            response = json.load(f)
        self.assertEqual(response["type"], "result")
        self.assertEqual(response["data"], 3)


class TestAsyncFunctionHandling(unittest.TestCase):
    """Tests for async function handling in handle_call."""

    def setUp(self):
        self.temp_dir = tempfile.mkdtemp()
        self.request_file = os.path.join(self.temp_dir, "request.json")
        self.response_file = os.path.join(self.temp_dir, "response.json")

        self.mock_module = type("MockModule", (), {"__name__": "test_module"})()

        async def async_add(a, b):
            await asyncio.sleep(0.001)
            return a + b

        self.mock_module.async_add = async_add

    def tearDown(self):
        import shutil

        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_async_function_in_handle_call(self) -> None:
        """Tests that async functions work correctly through handle_call."""
        with open(self.request_file, "w") as f:
            json.dump([5, 7], f)

        with patch("sys.stdout", io.StringIO()):
            persistent_wrapper.handle_call(
                f"CALL|async_add|{self.request_file}|{self.response_file}",
                self.mock_module,
                "default_func",
            )

        with open(self.response_file) as f:
            response = json.load(f)

        self.assertEqual(response["type"], "result")
        self.assertEqual(response["data"], 12)


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