fix: added a test for prompts (#1197)

explodinggradients · Aug 14, 2024 · b988f53 · b988f53
1 parent a77c475
commit b988f53
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/src/ragas/llms/prompt.py b/src/ragas/llms/prompt.py
@@ -269,7 +269,7 @@ def get_all_keys(nested_json):
 
         return self
 
-    def save(self, cache_dir: t.Optional[str] = None) -> None:
+    def save(self, cache_dir: t.Optional[str] = None):
         cache_dir = cache_dir if cache_dir else get_cache_dir()
         cache_dir = os.path.join(cache_dir, self.language)
         if not os.path.exists(cache_dir):

diff --git a/tests/unit/llms/test_prompt.py b/tests/unit/llms/test_prompt.py
@@ -121,3 +121,12 @@ def test_prompt_object_names():
                     obj.name not in prompt_object_names
                 ), f"Duplicate prompt name: {obj.name}"
                 prompt_object_names.append(obj.name)
+
+
+def test_save_and_load(tmp_path):
+    for testcase in TESTCASES:
+        prompt = Prompt(**testcase)
+        prompt.save(tmp_path)
+        loaded_prompt = prompt._load(prompt.language, prompt.name, tmp_path)
+
+        assert prompt == loaded_prompt