From 95ee4cf300c4d69b36c6f3a75fcf0ea9b1315d96 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20And=C3=A9n?= <janden@kth.se>
Date: Tue, 27 Aug 2024 20:53:08 +0200
Subject: [PATCH] Fix cufinufft fallback (#536)

* py: fix bug in cufinufft library fallback

* py: add tests for fallback library loading

* py: skip fallback tests in CI

These tests checks that the fallback mechanism is working properly. In
other words, if there is no bundled library, it should look in the
system path (e.g. `LD_LIBRARY_PATH` on Unix systems) and try to load it
from there.

For some reason, these tests work locally, but in CI, the library is
imported as usual, even with the patched function. Perhaps this has
something to do with the fact that the module is already loaded, so
`import finufft` is a no-op (this is evidenced by the fact that any
`print` statement in `finufft._finufft` is not detected in the output).
---
 Jenkinsfile                              |  2 +-
 python/cufinufft/cufinufft/_cufinufft.py |  1 +
 python/cufinufft/tests/test_fallback.py  | 25 ++++++++++++++++++++++++
 python/finufft/pyproject.toml            |  2 +-
 python/finufft/test/test_fallback.py     | 20 +++++++++++++++++++
 5 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 python/cufinufft/tests/test_fallback.py
 create mode 100644 python/finufft/test/test_fallback.py

diff --git a/Jenkinsfile b/Jenkinsfile
index ac7fd8025..daf1daf47 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -59,7 +59,7 @@ pipeline {
       source $HOME/bin/activate
       python3 -m pip install --no-cache-dir --upgrade pycuda cupy-cuda112 numba
       python3 -m pip install --no-cache-dir torch==1.12.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
-      python3 -m pip install --no-cache-dir pytest
+      python3 -m pip install --no-cache-dir pytest pytest-mock
       python -c "from numba import cuda; cuda.cudadrv.libs.test()"
       python3 -m pytest --framework=pycuda python/cufinufft
       python3 -m pytest --framework=numba python/cufinufft
diff --git a/python/cufinufft/cufinufft/_cufinufft.py b/python/cufinufft/cufinufft/_cufinufft.py
index 4c62fe947..ee51f0490 100644
--- a/python/cufinufft/cufinufft/_cufinufft.py
+++ b/python/cufinufft/cufinufft/_cufinufft.py
@@ -11,6 +11,7 @@
 import importlib.util
 import pathlib
 import numpy as np
+from ctypes.util import find_library
 
 from ctypes import c_double
 from ctypes import c_int
diff --git a/python/cufinufft/tests/test_fallback.py b/python/cufinufft/tests/test_fallback.py
new file mode 100644
index 000000000..903561551
--- /dev/null
+++ b/python/cufinufft/tests/test_fallback.py
@@ -0,0 +1,25 @@
+import pytest
+
+import numpy as np
+from ctypes.util import find_library
+
+
+# Check to make sure the fallback mechanism works if there is no bundled
+# dynamic library.
+@pytest.mark.skip(reason="Patching seems to fail in CI")
+def test_fallback(mocker):
+    def fake_load_library(lib_name, path):
+        if lib_name in ["libcufinufft", "cufinufft"]:
+            raise OSError()
+        else:
+            return np.ctypeslib.load_library(lib_name, path)
+
+    # Block out the bundled library.
+    mocker.patch("numpy.ctypeslib.load_library", fake_load_library)
+
+    # Make sure an error is raised if no system library is found.
+    if find_library("cufinufft") is None:
+        with pytest.raises(ImportError, match="suitable cufinufft"):
+            import cufinufft
+    else:
+        import cufinufft
diff --git a/python/finufft/pyproject.toml b/python/finufft/pyproject.toml
index f75ed4445..cf0309d3b 100644
--- a/python/finufft/pyproject.toml
+++ b/python/finufft/pyproject.toml
@@ -59,7 +59,7 @@ input = "finufft/__init__.py"
 build-verbosity = 1
 # Not building for PyPy and musllinux for now.
 skip = "pp* *musllinux*"
-test-requires = "pytest"
+test-requires = ["pytest", "pytest-mock"]
 test-command = "pytest {project}/python/finufft/test"
 
 [tool.cibuildwheel.linux]
diff --git a/python/finufft/test/test_fallback.py b/python/finufft/test/test_fallback.py
new file mode 100644
index 000000000..5b1e529f1
--- /dev/null
+++ b/python/finufft/test/test_fallback.py
@@ -0,0 +1,20 @@
+import pytest
+
+import numpy as np
+from ctypes.util import find_library
+
+@pytest.mark.skip(reason="Patching seems to fail in CI")
+def test_fallback(mocker):
+    def fake_load_library(lib_name, path):
+        if lib_name in ["libfinufft", "finufft"]:
+            raise OSError()
+        else:
+            return np.ctypeslib.load_library(lib_name, path)
+
+    mocker.patch("numpy.ctypeslib.load_library", fake_load_library)
+
+    if find_library("finufft") is None:
+        with pytest.raises(ImportError, match="suitable finufft"):
+            import finufft
+    else:
+        import finufft