test(body test): implement

ovuruska · May 4, 2024 · ccee46b · ccee46b
1 parent c1975e3
commit ccee46b
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py
@@ -1,3 +1,4 @@
+import json
 import unittest
 from unittest.mock import patch
 
@@ -24,9 +25,9 @@ def test_generate(self, mock_post):
         called_args, called_kwargs = mock_post.call_args
         url = called_args[0]
         header = called_kwargs["headers"]
-        body = called_kwargs["data"]
+        data = called_kwargs["data"]
 
         self.assertEqual(url, f"https://api.deepinfra.com/v1/inference/{model_name}")
-        self.assertEqual(body, '{"text": "Hello, World!"}')
+        self.assertEqual(data, json.dumps(body))
         self.assertEqual(response.embeddings, [1, 2, 3])
         self.assertEqual(header["Authorization"], f"Bearer {api_key}")
diff --git a/tests/test_text_generation.py b/tests/test_text_generation.py
@@ -1,3 +1,4 @@
+import json
 import unittest
 from unittest.mock import patch
 
@@ -25,8 +26,10 @@ def test_generate(self, mock_post):
 
         called_args, called_kwargs = mock_post.call_args
         url = called_args[0]
+        data = called_kwargs["data"]
         header = called_kwargs["headers"]
         self.assertEqual(url, f"https://api.deepinfra.com/v1/inference/{model_name}")
 
         self.assertEqual(response.results, [])
         self.assertEqual(header["Authorization"], f"Bearer {api_key}")
+        self.assertEqual(data, json.dumps(body))
diff --git a/tests/test_text_to_image.py b/tests/test_text_to_image.py
@@ -1,3 +1,4 @@
+import json
 import unittest
 from unittest.mock import patch
 
@@ -30,7 +31,9 @@ def test_generate(self, mock_post):
         called_args, called_kwargs = mock_post.call_args
         url = called_args[0]
         header = called_kwargs["headers"]
+        data = called_kwargs["data"]
         self.assertEqual(url, f"https://api.deepinfra.com/v1/inference/{model_name}")
 
         self.assertEqual(response.images, images)
         self.assertEqual(header["Authorization"], f"Bearer {api_key}")
+        self.assertEqual(data, json.dumps(body))