tests: add unit tests for Nvidia integration (#3162)

Author: anakin87

integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py

@@ -253,7 +253,7 @@ def run(self, documents: list[Document]) -> dict[str, list[Document] | dict[str,
         if not self._initialized:
             self.warm_up()
 
-        elif not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
+        if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
             msg = (
                 "NvidiaDocumentEmbedder expects a list of Documents as input."
                 "In case you want to embed a string, please use the NvidiaTextEmbedder."

integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py

@@ -205,13 +205,13 @@ def run(self, text: str) -> dict[str, list[float] | dict[str, Any]]:
         if not self._initialized:
             self.warm_up()
 
-        elif not isinstance(text, str):
+        if not isinstance(text, str):
             msg = (
                 "NvidiaTextEmbedder expects a string as an input."
                 "In case you want to embed a list of Documents, please use the NvidiaDocumentEmbedder."
             )
             raise TypeError(msg)
-        elif not text:
+        if not text:
             msg = "Cannot embed an empty string."
             raise ValueError(msg)
 

integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py

@@ -132,7 +132,7 @@ def warm_up(self) -> None:
 
         if not self.is_hosted and not self._model:
             if self.backend.model:
-                self.model = self.backend.model
+                self._model = self.backend.model
             else:
                 self.default_model()
 

integrations/nvidia/tests/test_document_embedder.py

@@ -354,6 +354,13 @@ def test_run_wrong_input_format(self):
         with pytest.raises(TypeError, match="NvidiaDocumentEmbedder expects a list of Documents as input"):
             embedder.run(documents=list_integers_input)
 
+    def test_run_validates_input_without_prior_warm_up(self):
+        api_key = Secret.from_token("fake-api-key")
+        embedder = NvidiaDocumentEmbedder("nvidia/nv-embedqa-e5-v5", api_key=api_key)
+
+        with pytest.raises(TypeError, match="NvidiaDocumentEmbedder expects a list of Documents as input"):
+            embedder.run(documents="text")
+
     def test_run_empty_document(self, caplog):
         model = "nvidia/nv-embedqa-e5-v5"
         api_key = Secret.from_token("fake-api-key")

integrations/nvidia/tests/test_generator.py

@@ -230,3 +230,73 @@ def test_hosted_nim_without_key(self):
         )
         with pytest.raises(ValueError):
             generator1.warm_up()
+
+    @pytest.mark.usefixtures("mock_local_models")
+    def test_warm_up_falls_back_to_default_model(self, monkeypatch):
+        monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+        generator = NvidiaGenerator(api_url="http://localhost:8080/v1")
+
+        with pytest.warns(UserWarning, match="Default model is set as:"):
+            generator.warm_up()
+
+        assert generator._model == "model1"
+        assert generator.backend.model == "model1"
+        assert generator.to_dict()["init_parameters"]["model"] == "model1"
+
+    def test_default_model_raises_when_no_valid_models(self, monkeypatch, requests_mock):
+        monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+        requests_mock.get(
+            "http://localhost:8080/v1/models",
+            json={"data": [{"id": "derived-model", "object": "model", "root": "base-model"}]},
+        )
+        generator = NvidiaGenerator(api_url="http://localhost:8080/v1")
+
+        with pytest.raises(ValueError, match="No locally hosted model was found"):
+            generator.warm_up()
+
+    def test_warm_up_is_idempotent(self, monkeypatch):
+        monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+        generator = NvidiaGenerator("meta/llama3-8b-instruct")
+        generator.warm_up()
+        backend = generator.backend
+        generator.warm_up()
+        assert generator.backend is backend
+
+    def test_available_models_without_backend(self, monkeypatch):
+        monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+        generator = NvidiaGenerator("meta/llama3-8b-instruct")
+        assert generator.available_models == []
+
+    @pytest.mark.usefixtures("mock_local_models")
+    def test_available_models_with_backend(self, monkeypatch):
+        monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+        generator = NvidiaGenerator(model="model1", api_url="http://localhost:8080/v1")
+        generator.warm_up()
+        models = generator.available_models
+        assert len(models) == 1
+        assert models[0].id == "model1"
+
+    def test_from_dict(self, monkeypatch):
+        monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+        data = {
+            "type": "haystack_integrations.components.generators.nvidia.generator.NvidiaGenerator",
+            "init_parameters": {
+                "api_key": {"env_vars": ["NVIDIA_API_KEY"], "strict": True, "type": "env_var"},
+                "api_url": "https://my.url.com/v1",
+                "model": "meta/llama3-8b-instruct",
+                "model_arguments": {"temperature": 0.5},
+            },
+        }
+        generator = NvidiaGenerator.from_dict(data)
+        assert generator._model == "meta/llama3-8b-instruct"
+        assert generator.api_url == "https://my.url.com/v1"
+        assert generator._model_arguments == {"temperature": 0.5}
+
+    def test_run(self, monkeypatch, mock_local_chat_completion):  # noqa: ARG002
+        monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+        generator = NvidiaGenerator(model="model1", api_url="http://localhost:8080/v1")
+
+        result = generator.run(prompt="What is the answer?")
+
+        assert result["replies"] == ["Hello!", "How are you?"]
+        assert len(result["meta"]) == 2

integrations/nvidia/tests/test_nim_backend.py

@@ -233,6 +233,36 @@ def test_models(self, monkeypatch):
                 timeout=60.0,
             )
 
+    def test_embed_raises_on_http_error(self, monkeypatch):
+        error_response = requests.Response()
+        error_response.status_code = 500
+        error_response._content = b"server exploded"
+        with patch("requests.sessions.Session.post", return_value=error_response):
+            monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+            backend = NimBackend(model="nvidia/nv-embedqa-e5-v5", api_url=DEFAULT_API_URL, client="NvidiaTextEmbedder")
+            with pytest.raises(ValueError, match="Failed to query embedding endpoint"):
+                backend.embed(texts=["a"])
+
+    def test_generate_raises_on_http_error(self, monkeypatch):
+        error_response = requests.Response()
+        error_response.status_code = 500
+        error_response._content = b"server exploded"
+        with patch("requests.sessions.Session.post", return_value=error_response):
+            monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+            backend = NimBackend(model="meta/llama3-8b-instruct", api_url=DEFAULT_API_URL, client="NvidiaGenerator")
+            with pytest.raises(ValueError, match="Failed to query chat completion endpoint"):
+                backend.generate(prompt="hi")
+
+    def test_models_raises_when_empty(self, monkeypatch):
+        empty_response = requests.Response()
+        empty_response.status_code = 200
+        empty_response._content = json.dumps({"data": []}).encode()
+        with patch("requests.sessions.Session.get", return_value=empty_response):
+            monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+            backend = NimBackend(model="custom-model", api_url="http://localhost:8000")
+            with pytest.raises(ValueError, match="No hosted model were found"):
+                backend.models()
+
     def test_rank(self, monkeypatch):
         with patch("requests.sessions.Session.post", side_effect=mock_rank_post_response) as mock_post:
             monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
@@ -258,3 +288,27 @@ def test_rank(self, monkeypatch):
                 },
                 timeout=60.0,
             )
+
+    def test_rank_raises_on_http_error(self, monkeypatch):
+        error_response = requests.Response()
+        error_response.status_code = 500
+        error_response._content = b"server exploded"
+        with patch("requests.sessions.Session.post", return_value=error_response):
+            monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+            backend = NimBackend(
+                model="nvidia/llama-3.2-nv-rerankqa-1b-v2", api_url=DEFAULT_API_URL, client="NvidiaRanker"
+            )
+            with pytest.raises(ValueError, match="Failed to rank endpoint"):
+                backend.rank(query_text="q", document_texts=["a"])
+
+    def test_rank_raises_when_rankings_missing(self, monkeypatch):
+        response = requests.Response()
+        response.status_code = 200
+        response._content = json.dumps({"unexpected": "payload"}).encode()
+        with patch("requests.sessions.Session.post", return_value=response):
+            monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+            backend = NimBackend(
+                model="nvidia/llama-3.2-nv-rerankqa-1b-v2", api_url=DEFAULT_API_URL, client="NvidiaRanker"
+            )
+            with pytest.raises(ValueError, match="Expected 'rankings' in response"):
+                backend.rank(query_text="q", document_texts=["a"])

integrations/nvidia/tests/test_ranker.py

@@ -330,6 +330,20 @@ def test_setting_timeout_env(self, monkeypatch):
         client.warm_up()
         assert client.backend.timeout == 45.0
 
+    def test_run_on_empty_list(self, monkeypatch):
+        monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+        client = NvidiaRanker()
+        client.warm_up()
+        assert client.run(query="q", documents=[]) == {"documents": []}
+
+    def test_run_without_prior_warm_up(self, requests_mock, monkeypatch):
+        monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
+        requests_mock.post(re.compile(r".*ranking"), json={"rankings": [{"index": 0, "logit": 1.0}]})
+        client = NvidiaRanker()
+        result = client.run(query="q", documents=[Document(content="doc")])
+        assert client._initialized is True
+        assert len(result["documents"]) == 1
+
     def test_prepare_texts_to_embed_w_metadata(self):
         documents = [
             Document(content=f"document number {i}:\ncontent", meta={"meta_field": f"meta_value {i}"}) for i in range(5)

integrations/nvidia/tests/test_text_embedder.py

@@ -183,6 +183,20 @@ def test_run_empty_string(self):
         with pytest.raises(ValueError, match="empty string"):
             embedder.run(text="")
 
+    def test_run_validates_input_without_prior_warm_up(self):
+        api_key = Secret.from_token("fake-api-key")
+        embedder = NvidiaTextEmbedder("nvidia/nv-embedqa-e5-v5", api_key=api_key)
+
+        with pytest.raises(TypeError, match="NvidiaTextEmbedder expects a string as an input"):
+            embedder.run(text=[1, 2, 3])
+
+    def test_run_rejects_empty_string_without_prior_warm_up(self):
+        api_key = Secret.from_token("fake-api-key")
+        embedder = NvidiaTextEmbedder("nvidia/nv-embedqa-e5-v5", api_key=api_key)
+
+        with pytest.raises(ValueError, match="empty string"):
+            embedder.run(text="")
+
     def test_setting_timeout(self, monkeypatch):
         monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
         embedder = NvidiaTextEmbedder(timeout=10.0)

integrations/nvidia/tests/test_utils.py

@@ -5,7 +5,12 @@
 import pytest
 
 from haystack_integrations.utils.nvidia import Client, is_hosted
-from haystack_integrations.utils.nvidia.models import CHAT_MODEL_TABLE, EMBEDDING_MODEL_TABLE, RANKING_MODEL_TABLE
+from haystack_integrations.utils.nvidia.models import (
+    CHAT_MODEL_TABLE,
+    EMBEDDING_MODEL_TABLE,
+    RANKING_MODEL_TABLE,
+    Model,
+)
 from haystack_integrations.utils.nvidia.utils import (
     determine_model,
     lookup_model,
@@ -110,3 +115,39 @@ def test_validate_hosted_model_with_client() -> None:
     model = validate_hosted_model("meta/codellama-70b", Client.NVIDIA_GENERATOR)
     assert model is not None
     assert model.client == Client.NVIDIA_GENERATOR
+
+
+# Model
+def test_model_hash_uses_id() -> None:
+    assert hash(Model(id="foo")) == hash("foo")
+
+
+def test_model_validate_with_enum_client() -> None:
+    model = Model(id="foo", model_type="chat", client=Client.NVIDIA_GENERATOR)
+    assert model.validate() == hash("foo")
+
+
+def test_model_validate_with_string_client() -> None:
+    model = Model(id="foo", model_type="chat", client="NvidiaGenerator")
+    assert model.validate() == hash("foo")
+
+
+def test_model_validate_raises_on_incompatible_type() -> None:
+    model = Model(id="foo", model_type="embedding", client=Client.NVIDIA_GENERATOR)
+    with pytest.raises(ValueError, match="not supported by client"):
+        model.validate()
+
+
+def test_model_validate_without_client() -> None:
+    model = Model(id="foo", model_type="chat")
+    assert model.validate() == hash("foo")
+
+
+# Client
+def test_client_from_str_invalid() -> None:
+    with pytest.raises(ValueError, match="Unknown client"):
+        Client.from_str("NotARealClient")
+
+
+def test_client_str_returns_value() -> None:
+    assert str(Client.NVIDIA_GENERATOR) == "NvidiaGenerator"