From bd7a60bb234ee1766b9362004e40eb365b17ddc3 Mon Sep 17 00:00:00 2001
From: Um Changyong <changyong.um@sfa.co.kr>
Date: Thu, 26 Dec 2024 14:07:44 +0900
Subject: [PATCH] feat: add tests

---
 tests/autorag/embedding/test_base.py | 39 ++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)
 create mode 100644 tests/autorag/embedding/test_base.py

diff --git a/tests/autorag/embedding/test_base.py b/tests/autorag/embedding/test_base.py
new file mode 100644
index 000000000..1feb9a748
--- /dev/null
+++ b/tests/autorag/embedding/test_base.py
@@ -0,0 +1,39 @@
+import pytest
+from autorag.embedding.base import EmbeddingModel
+
+
+def test_load_embedding_model():
+	# Test loading a supported embedding model
+	embedding = EmbeddingModel.load("mock")
+	assert embedding is not None
+
+	# Test loading an unsupported embedding model
+	with pytest.raises(
+		ValueError, match="Embedding model 'unsupported_model' is not supported"
+	):
+		EmbeddingModel.load("unsupported_model")
+
+
+def test_load_embedding_model_from_dict():
+	# Test loading with missing keys
+	with pytest.raises(
+		ValueError, match="Both 'type' and 'model_name' must be provided"
+	):
+		EmbeddingModel.load_from_dict([{"type": "openai"}])
+
+	# Test loading with unsupported type
+	with pytest.raises(
+		ValueError, match="Embedding model type 'unsupported_type' is not supported"
+	):
+		EmbeddingModel.load_from_dict(
+			[{"type": "unsupported_type", "model_name": "some-model"}]
+		)
+
+	# Test loading with multiple items
+	with pytest.raises(ValueError, match="Only one embedding model is supported"):
+		EmbeddingModel.load_from_dict(
+			[
+				{"type": "openai", "model_name": "text-embedding-ada-002"},
+				{"type": "huggingface", "model_name": "BAAI/bge-small-en-v1.5"},
+			]
+		)