test(cuda.core): align nvJitLink version cache coverage

Author: isVoid

cuda_core/tests/test_optional_dependency_imports.py

@@ -18,12 +18,16 @@ def restore_optional_import_state():
     saved_nvvm_attempted = _program._nvvm_import_attempted
     saved_driver = _linker._driver
     saved_inited = _linker._inited
+    saved_nvjitlink = _linker._nvjitlink
+    saved_nvjitlink_version = _linker._nvjitlink_version
     saved_use_nvjitlink = _linker._use_nvjitlink_backend
 
     _program._nvvm_module = None
     _program._nvvm_import_attempted = False
     _linker._driver = None
     _linker._inited = False
+    _linker._nvjitlink = None
+    _linker._nvjitlink_version = None
     _linker._use_nvjitlink_backend = None
 
     yield
@@ -32,6 +36,8 @@ def restore_optional_import_state():
     _program._nvvm_import_attempted = saved_nvvm_attempted
     _linker._driver = saved_driver
     _linker._inited = saved_inited
+    _linker._nvjitlink = saved_nvjitlink
+    _linker._nvjitlink_version = saved_nvjitlink_version
     _linker._use_nvjitlink_backend = saved_use_nvjitlink
 
 
@@ -168,12 +174,22 @@ def fake__optional_cuda_import(modname, probe_function=None):
     assert _linker._use_nvjitlink_backend is False
 
 
-@pytest.mark.agent_authored(model="grok-4.5")
+@pytest.mark.agent_authored(model="gpt-5.6")
 def test_decide_nvjitlink_or_driver_selects_nvjitlink_when_version_symbol_present(monkeypatch):
+    version_calls = 0
+
+    class FakeModule:
+        def version(self):
+            nonlocal version_calls
+            version_calls += 1
+            return (13, 4)
+
+    nvjitlink_module = FakeModule()
+
     def fake__optional_cuda_import(modname, probe_function=None):
         assert modname == "cuda.bindings.nvjitlink"
         assert probe_function is None
-        return object()
+        return nvjitlink_module
 
     monkeypatch.setattr(_linker, "_optional_cuda_import", fake__optional_cuda_import)
     monkeypatch.setattr(_linker, "_nvjitlink_has_version_symbol", lambda _nvjitlink: True)
@@ -182,21 +198,24 @@ def fake__optional_cuda_import(modname, probe_function=None):
 
     assert use_driver_backend is False
     assert _linker._use_nvjitlink_backend is True
+    assert _linker._nvjitlink is nvjitlink_module
+    assert _linker._nvjitlink_version == (13, 4)
+    assert version_calls == 1
 
 
-@pytest.mark.agent_authored(model="grok-4.5")
-def test_decide_nvjitlink_or_driver_does_not_call_version(monkeypatch):
-    """Regression guard for #2408: must not call module.version()."""
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_decide_nvjitlink_or_driver_does_not_call_version_when_symbol_missing(monkeypatch):
+    """Regression guard for #2408: old nvJitLink must not call module.version()."""
     called = {"version": False, "inspect": False}
 
     class FakeModule:
         def version(self):
             called["version"] = True
-            raise AssertionError("module.version() must not be used for nvJitLink probing")
+            raise AssertionError("module.version() must not be called when its symbol is missing")
 
     def fake_has_version(_nvjitlink):
         called["inspect"] = True
-        return True
+        return False
 
     def fake__optional_cuda_import(modname, probe_function=None):
         assert modname == "cuda.bindings.nvjitlink"
@@ -206,6 +225,9 @@ def fake__optional_cuda_import(modname, probe_function=None):
     monkeypatch.setattr(_linker, "_optional_cuda_import", fake__optional_cuda_import)
     monkeypatch.setattr(_linker, "_nvjitlink_has_version_symbol", fake_has_version)
 
-    assert _linker._decide_nvjitlink_or_driver() is False
+    with pytest.warns(RuntimeWarning, match="too old \\(<12.3\\)"):
+        assert _linker._decide_nvjitlink_or_driver() is True
     assert called["inspect"] is True
     assert called["version"] is False
+    assert _linker._nvjitlink is None
+    assert _linker._nvjitlink_version is None