From d74a6b8983ea20d0b8f33435ccbe5f56e8f5019f Mon Sep 17 00:00:00 2001
From: Marnix Kraus <marnix.kraus@deltares.nl>
Date: Mon, 7 Oct 2024 08:39:39 +0200
Subject: [PATCH] Test singledispatch in isolated subprocess

---
 flopy4/singledispatch/plot.py     |  3 +-
 flopy4/singledispatch/plot_int.py |  5 ++-
 test/test_singledispatch.py       | 54 +++++++++++++++++++++++++++++--
 3 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/flopy4/singledispatch/plot.py b/flopy4/singledispatch/plot.py
index bb7879e..31933b4 100644
--- a/flopy4/singledispatch/plot.py
+++ b/flopy4/singledispatch/plot.py
@@ -1,8 +1,9 @@
 from functools import singledispatch
+from typing import Any
 
 
 @singledispatch
-def plot(obj, **kwargs):
+def plot(obj, **kwargs) -> Any:
     raise NotImplementedError(
         "plot method not implemented for type {}".format(type(obj))
     )
diff --git a/flopy4/singledispatch/plot_int.py b/flopy4/singledispatch/plot_int.py
index 6da9e6a..40ab777 100644
--- a/flopy4/singledispatch/plot_int.py
+++ b/flopy4/singledispatch/plot_int.py
@@ -1,6 +1,9 @@
+from typing import Any
+
 from flopy4.singledispatch.plot import plot
 
 
 @plot.register
-def _(v: int, **kwargs):
+def _(v: int, **kwargs) -> Any:
     print(f"Plotting a model with kwargs: {kwargs}")
+    return v
diff --git a/test/test_singledispatch.py b/test/test_singledispatch.py
index 2a98597..6e5337b 100644
--- a/test/test_singledispatch.py
+++ b/test/test_singledispatch.py
@@ -1,12 +1,62 @@
+import ast
+import inspect
+import subprocess
+import sys
 from importlib.metadata import entry_points
 
+import pytest
+
+from flopy4.singledispatch.plot import plot
+
+
+def get_function_body(func):
+    source = inspect.getsource(func)
+    parsed = ast.parse(source)
+    for node in ast.walk(parsed):
+        if isinstance(node, ast.FunctionDef):
+            return ast.get_source_segment(source, node.body[0])
+    raise ValueError("Function body not found")
+
+
+def run_test_in_subprocess(test_func):
+    def wrapper():
+        test_func_source = get_function_body(test_func)
+        test_code = f"""
+import pytest
+from importlib.metadata import entry_points
 from flopy4.singledispatch.plot import plot
 
+{test_func_source}
+
+"""
+        result = subprocess.run(
+            [sys.executable, "-c", test_code], capture_output=True, text=True
+        )
+        if result.returncode != 0:
+            print(result.stdout)
+            print(result.stderr)
+        assert result.returncode == 0, f"Test failed: {test_func.__name__}"
+
+    return wrapper
 
+
+@run_test_in_subprocess
 def test_register_singledispatch_with_entrypoints():
     eps = entry_points(group="flopy4", name="plot")
     for ep in eps:
-        _ = ep.load()
+        ep.load()
 
     # should not throw an error, because plot_int was loaded via entry points
-    plot(5)
+    return_val = plot(5)
+    assert return_val == 5
+    with pytest.raises(NotImplementedError):
+        plot("five")
+
+
+@run_test_in_subprocess
+def test_register_singledispatch_without_entrypoints():
+    # should throw an error, because plot_int was not loaded via entry points
+    with pytest.raises(NotImplementedError):
+        plot(5)
+    with pytest.raises(NotImplementedError):
+        plot("five")