From 60d39f5f7f175f94b2511b221ee2fd1760eacb9e Mon Sep 17 00:00:00 2001
From: dangotbanned <125183946+dangotbanned@users.noreply.github.com>
Date: Mon, 11 Nov 2024 16:40:12 +0000
Subject: [PATCH] test: Adds tests for missing dependencies

---
 altair/datasets/_readers.py | 14 ++++++++++-
 tests/test_datasets.py      | 48 +++++++++++++++++++++++++++++++++++--
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/altair/datasets/_readers.py b/altair/datasets/_readers.py
index b2f41af89..20b308aed 100644
--- a/altair/datasets/_readers.py
+++ b/altair/datasets/_readers.py
@@ -226,7 +226,19 @@ def _import(self, name: str, /) -> Any:
         if spec := find_spec(name):
             return import_module(spec.name)
         else:
-            msg = f"{type(self).__name__!r} requires missing dependency {name!r}."
+            reqs = _requirements(self._name)  # type: ignore[call-overload]
+            if isinstance(reqs, tuple):
+                depends = ", ".join(f"{req!r}" for req in reqs) + " packages"
+            else:
+                depends = f"{reqs!r} package"
+
+            msg = (
+                f"Backend {self._name!r} requires the {depends}, but {name!r} could not be found.\n"
+                f"This can be installed with pip using:\n"
+                f"    pip install {name}\n"
+                f"Or with conda using:\n"
+                f"    conda install -c conda-forge {name}"
+            )
             raise ModuleNotFoundError(msg, name=name)
 
     def __repr__(self) -> str:
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 7a4ab51f1..de932137f 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import re
+import sys
 from importlib.util import find_spec
 from typing import TYPE_CHECKING
 
@@ -13,8 +14,12 @@
 from tests import skip_requires_pyarrow
 
 if TYPE_CHECKING:
+    from typing import Literal
+
     from altair.datasets._readers import _Backend
 
+CACHE_ENV_VAR: Literal["ALTAIR_DATASETS_DIR"] = "ALTAIR_DATASETS_DIR"
+
 
 requires_pyarrow = skip_requires_pyarrow()
 
@@ -58,10 +63,49 @@ def test_loader_url(backend: _Backend) -> None:
 
 
 @backends
-def test_loader_call(backend: _Backend) -> None:
+def test_loader_call(backend: _Backend, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv(CACHE_ENV_VAR, raising=False)
+
     data = Loader.with_backend(backend)
-    data.cache_dir = ""  # type: ignore[assignment]
     frame = data("stocks", ".csv")
     assert is_into_dataframe(frame)
     nw_frame = nw.from_native(frame)
     assert set(nw_frame.columns) == {"symbol", "date", "price"}
+
+
+@backends
+def test_missing_dependency_single(
+    backend: _Backend, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    if backend in {"polars[pyarrow]", "pandas[pyarrow]"}:
+        pytest.skip("Testing single dependency backends only")
+
+    monkeypatch.setitem(sys.modules, backend, None)
+
+    with pytest.raises(
+        ModuleNotFoundError,
+        match=re.compile(
+            rf"{backend}.+requires.+{backend}.+but.+{backend}.+not.+found.+pip install {backend}",
+            flags=re.DOTALL,
+        ),
+    ):
+        Loader.with_backend(backend)
+
+
+@pytest.mark.parametrize("backend", ["polars[pyarrow]", "pandas[pyarrow]"])
+@skip_requires_pyarrow
+def test_missing_dependency_multi(
+    backend: _Backend, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    secondary = "pyarrow"
+    primary = backend.removesuffix(f"[{secondary}]")
+    monkeypatch.setitem(sys.modules, secondary, None)
+
+    with pytest.raises(
+        ModuleNotFoundError,
+        match=re.compile(
+            rf"{re.escape(backend)}.+requires.+'{primary}', '{secondary}'.+but.+{secondary}.+not.+found.+pip install {secondary}",
+            flags=re.DOTALL,
+        ),
+    ):
+        Loader.with_backend(backend)