From 29a4cc045e1cd6bdb603ae984b4ea3e89fed56f1 Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Fri, 8 May 2026 07:41:34 -0700
Subject: [PATCH] Test

---
 src/paperless_ai/tests/test_embedding.py | 28 ++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/src/paperless_ai/tests/test_embedding.py b/src/paperless_ai/tests/test_embedding.py
index e4e80cdf1..b595b68cd 100644
--- a/src/paperless_ai/tests/test_embedding.py
+++ b/src/paperless_ai/tests/test_embedding.py
@@ -100,6 +100,34 @@ def test_get_embedding_model_huggingface(mock_ai_config):
         assert model == MockHuggingFaceEmbedding.return_value
 
 
+def test_get_embedding_model_ollama(mock_ai_config):
+    mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.OLLAMA
+    mock_ai_config.return_value.llm_embedding_model = "embeddinggemma"
+    mock_ai_config.return_value.llm_endpoint = "http://test-url"
+
+    with patch(
+        "llama_index.embeddings.ollama.OllamaEmbedding",
+    ) as MockOllamaEmbedding:
+        model = get_embedding_model()
+        MockOllamaEmbedding.assert_called_once_with(
+            model_name="embeddinggemma",
+            base_url="http://test-url",
+        )
+        assert model == MockOllamaEmbedding.return_value
+
+
+def test_get_embedding_model_ollama_blocks_internal_endpoint_when_disallowed(
+    mock_ai_config,
+):
+    mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.OLLAMA
+    mock_ai_config.return_value.llm_embedding_model = "embeddinggemma"
+    mock_ai_config.return_value.llm_endpoint = "http://127.0.0.1:11434"
+    mock_ai_config.return_value.llm_allow_internal_endpoints = False
+
+    with pytest.raises(ValueError, match="non-public address"):
+        get_embedding_model()
+
+
 def test_get_embedding_model_invalid_backend(mock_ai_config):
     mock_ai_config.return_value.llm_embedding_backend = "INVALID_BACKEND"