From af45b5a11a8a393ed404be3ce0ecac14338155c7 Mon Sep 17 00:00:00 2001
From: Vladimir Repin <32306715+mezotaken@users.noreply.github.com>
Date: Sat, 29 Oct 2022 18:26:28 +0300
Subject: [PATCH] Testing with API added

---
 run_tests.bat        | 15 +++++++++
 test/__init__.py     |  0
 test/server_poll.py  | 17 ++++++++++
 test/txt2img_test.py | 80 ++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 112 insertions(+)
 create mode 100644 run_tests.bat
 create mode 100644 test/__init__.py
 create mode 100644 test/server_poll.py
 create mode 100644 test/txt2img_test.py

diff --git a/run_tests.bat b/run_tests.bat
new file mode 100644
index 00000000..3a63f034
--- /dev/null
+++ b/run_tests.bat
@@ -0,0 +1,15 @@
+@echo off
+set ERROR_REPORTING=FALSE
+set COMMANDLINE_ARGS= --api
+echo Launching SDWebUI...
+start "SDWebUITest" webui.bat
+
+if not defined PYTHON (set PYTHON=python)
+if not defined VENV_DIR (set VENV_DIR=venv)
+set PYTHON="%~dp0%VENV_DIR%\Scripts\Python.exe"
+%PYTHON% test/server_poll.py
+for /f "tokens=2 delims=," %%a in ('tasklist /v /fo csv ^| findstr /i "SDWebUITest"') do set "$PID=%%a"
+
+taskkill /PID %$PID% >nul 2>&1
+
+pause
diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/test/server_poll.py b/test/server_poll.py
new file mode 100644
index 00000000..8c0436f8
--- /dev/null
+++ b/test/server_poll.py
@@ -0,0 +1,17 @@
+import unittest
+import requests
+import time
+
+timeout_threshold = 240
+start_time = time.time()
+while time.time()-start_time < timeout_threshold:
+    try:
+        requests.head("http://localhost:7860/")
+        break
+    except requests.exceptions.ConnectionError:
+        pass
+if time.time()-start_time < timeout_threshold:
+    suite = unittest.TestLoader().discover('', pattern='*_test.py')
+    result = unittest.TextTestRunner(verbosity=2).run(suite)
+else:
+    print("Launch unsuccessful")
diff --git a/test/txt2img_test.py b/test/txt2img_test.py
new file mode 100644
index 00000000..9484fd99
--- /dev/null
+++ b/test/txt2img_test.py
@@ -0,0 +1,80 @@
+import unittest
+import requests
+import time
+
+url_txt2img = "http://localhost:7860/sdapi/v1/txt2img"
+simple_txt2img = {
+    "enable_hr": False,
+    "denoising_strength": 0,
+    "firstphase_width": 0,
+    "firstphase_height": 0,
+    "prompt": "example prompt",
+    "styles": [
+      ""
+    ],
+    "seed": -1,
+    "subseed": -1,
+    "subseed_strength": 0,
+    "seed_resize_from_h": -1,
+    "seed_resize_from_w": -1,
+    "batch_size": 1,
+    "n_iter": 1,
+    "steps": 5,
+    "cfg_scale": 7,
+    "width": 64,
+    "height": 64,
+    "restore_faces": False,
+    "tiling": False,
+    "negative_prompt": "",
+    "eta": 0,
+    "s_churn": 0,
+    "s_tmax": 0,
+    "s_tmin": 0,
+    "s_noise": 1,
+    "sampler_index": "Euler a"
+}
+
+class TestTxt2ImgWorking(unittest.TestCase):
+    def test_txt2img_simple_performed(self):
+        self.assertEqual(requests.post(url_txt2img, json=simple_txt2img).status_code, 200)
+
+    def test_txt2img_with_negative_prompt_performed(self):
+        params = simple_txt2img.copy()
+        params["negative_prompt"] = "example negative prompt"
+        self.assertEqual(requests.post(url_txt2img, json=params).status_code, 200)
+
+    def test_txt2img_not_square_image_performed(self):
+        params = simple_txt2img.copy()
+        params["height"] = 128
+        self.assertEqual(requests.post(url_txt2img, json=params).status_code, 200)
+
+    def test_txt2img_with_hrfix_performed(self):
+        params = simple_txt2img.copy()
+        params["enable_hr"] = True
+        self.assertEqual(requests.post(url_txt2img, json=params).status_code, 200)
+
+    def test_txt2img_with_restore_faces_performed(self):
+        params = simple_txt2img.copy()
+        params["restore_faces"] = True
+        self.assertEqual(requests.post(url_txt2img, json=params).status_code, 200)
+
+    def test_txt2img_with_tiling_faces_performed(self):
+        params = simple_txt2img.copy()
+        params["tiling"] = True
+        self.assertEqual(requests.post(url_txt2img, json=params).status_code, 200)
+
+    def test_txt2img_with_vanilla_sampler_performed(self):
+        params = simple_txt2img.copy()
+        params["sampler_index"] = "PLMS"
+        self.assertEqual(requests.post(url_txt2img, json=params).status_code, 200)
+
+    def test_txt2img_multiple_batches_performed(self):
+        params = simple_txt2img.copy()
+        params["n_iter"] = 2
+        self.assertEqual(requests.post(url_txt2img, json=params).status_code, 200)
+
+class TestTxt2ImgCorrectness(unittest.TestCase):
+    pass
+
+if __name__ == "__main__":
+    unittest.main()