From 59dfe0845d964868e92572c78a420b6d68c46ea4 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 30 Oct 2022 08:22:44 +0300
Subject: [PATCH] launch tests from launch.py with --tests commandline argument

---
 .gitignore           |  2 ++
 launch.py            | 19 +++++++++++++++++++
 run_tests.bat        | 15 ---------------
 test/extras_test.py  |  6 +++---
 test/img2img_test.py | 12 ++++++------
 test/server_poll.py  | 28 +++++++++++++++-------------
 test/txt2img_test.py |  8 +++++---
 7 files changed, 50 insertions(+), 40 deletions(-)
 delete mode 100644 run_tests.bat

diff --git a/.gitignore b/.gitignore
index 8fa05852a..ee53044c5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -29,3 +29,5 @@ notification.mp3
 /textual_inversion
 .vscode
 /extensions
+/test/stdout.txt
+/test/stderr.txt
diff --git a/launch.py b/launch.py
index 8affd4109..33f98343c 100644
--- a/launch.py
+++ b/launch.py
@@ -128,10 +128,12 @@ def prepare_enviroment():
     blip_commit_hash = os.environ.get('BLIP_COMMIT_HASH', "48211a1594f1321b00f14c9f7a5b4813144b2fb9")
 
     sys.argv += shlex.split(commandline_args)
+    test_argv = [x for x in sys.argv if x != '--tests']
 
     sys.argv, skip_torch_cuda_test = extract_arg(sys.argv, '--skip-torch-cuda-test')
     sys.argv, reinstall_xformers = extract_arg(sys.argv, '--reinstall-xformers')
     sys.argv, update_check = extract_arg(sys.argv, '--update-check')
+    sys.argv, run_tests = extract_arg(sys.argv, '--tests')
     xformers = '--xformers' in sys.argv
     deepdanbooru = '--deepdanbooru' in sys.argv
     ngrok = '--ngrok' in sys.argv
@@ -194,6 +196,23 @@ def prepare_enviroment():
         print("Exiting because of --exit argument")
         exit(0)
 
+    if run_tests:
+        tests(test_argv)
+        exit(0)
+
+
+def tests(argv):
+    print(f"Launching Web UI in another process for testing with arguments: {' '.join(argv[1:])}")
+
+    with open('test/stdout.txt', "w", encoding="utf8") as stdout, open('test/stderr.txt', "w", encoding="utf8") as stderr:
+        proc = subprocess.Popen([sys.executable, *argv], stdout=stdout, stderr=stderr)
+
+    import test.server_poll
+    test.server_poll.run_tests()
+
+    print(f"Stopping Web UI process with id {proc.pid}")
+    proc.kill()
+
 
 def start_webui():
     print(f"Launching Web UI with arguments: {' '.join(sys.argv[1:])}")
diff --git a/run_tests.bat b/run_tests.bat
deleted file mode 100644
index 3a63f0349..000000000
--- a/run_tests.bat
+++ /dev/null
@@ -1,15 +0,0 @@
-@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/extras_test.py b/test/extras_test.py
index 2e1764d96..9b8ce0f03 100644
--- a/test/extras_test.py
+++ b/test/extras_test.py
@@ -1,7 +1,5 @@
 import unittest
-import requests
-from gradio.processing_utils import encode_pil_to_base64
-from PIL import Image
+
 
 class TestExtrasWorking(unittest.TestCase):
     def setUp(self):
@@ -22,8 +20,10 @@ class TestExtrasWorking(unittest.TestCase):
             "image": ""
             }
 
+
 class TestExtrasCorrectness(unittest.TestCase):
     pass
 
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/img2img_test.py b/test/img2img_test.py
index 61e3e285e..012a95809 100644
--- a/test/img2img_test.py
+++ b/test/img2img_test.py
@@ -3,13 +3,12 @@ import requests
 from gradio.processing_utils import encode_pil_to_base64
 from PIL import Image
 
+
 class TestImg2ImgWorking(unittest.TestCase):
     def setUp(self):
         self.url_img2img = "http://localhost:7860/sdapi/v1/img2img"
         self.simple_img2img = {
-            "init_images": [
-                encode_pil_to_base64(Image.open(r"test/test_files/img2img_basic.png"))
-                ],
+            "init_images": [encode_pil_to_base64(Image.open(r"test/test_files/img2img_basic.png"))],
             "resize_mode": 0,
             "denoising_strength": 0.75,
             "mask": None,
@@ -19,9 +18,7 @@ class TestImg2ImgWorking(unittest.TestCase):
             "inpaint_full_res_padding": 0,
             "inpainting_mask_invert": 0,
             "prompt": "example prompt",
-            "styles": [
-                ""
-                ],
+            "styles": [],
             "seed": -1,
             "subseed": -1,
             "subseed_strength": 0,
@@ -45,6 +42,7 @@ class TestImg2ImgWorking(unittest.TestCase):
             "sampler_index": "Euler a",
             "include_init_images": False
             }
+
     def test_img2img_simple_performed(self):
         self.assertEqual(requests.post(self.url_img2img, json=self.simple_img2img).status_code, 200)
 
@@ -52,8 +50,10 @@ class TestImg2ImgWorking(unittest.TestCase):
         self.simple_img2img["mask"] = encode_pil_to_base64(Image.open(r"test/test_files/mask_basic.png"))
         self.assertEqual(requests.post(self.url_img2img, json=self.simple_img2img).status_code, 200)
 
+
 class TestImg2ImgCorrectness(unittest.TestCase):
     pass
 
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/server_poll.py b/test/server_poll.py
index 8c0436f87..eeefb7ebc 100644
--- a/test/server_poll.py
+++ b/test/server_poll.py
@@ -2,16 +2,18 @@ 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")
+
+def run_tests():
+    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
index ad27686a9..1936e07e2 100644
--- a/test/txt2img_test.py
+++ b/test/txt2img_test.py
@@ -1,6 +1,7 @@
 import unittest
 import requests
 
+
 class TestTxt2ImgWorking(unittest.TestCase):
     def setUp(self):
         self.url_txt2img = "http://localhost:7860/sdapi/v1/txt2img"
@@ -10,9 +11,7 @@ class TestTxt2ImgWorking(unittest.TestCase):
             "firstphase_width": 0,
             "firstphase_height": 0,
             "prompt": "example prompt",
-            "styles": [
-            ""
-            ],
+            "styles": [],
             "seed": -1,
             "subseed": -1,
             "subseed_strength": 0,
@@ -34,6 +33,7 @@ class TestTxt2ImgWorking(unittest.TestCase):
             "s_noise": 1,
             "sampler_index": "Euler a"
         }
+
     def test_txt2img_simple_performed(self):
         self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
 
@@ -65,8 +65,10 @@ class TestTxt2ImgWorking(unittest.TestCase):
         self.simple_txt2img["n_iter"] = 2
         self.assertEqual(requests.post(self.url_txt2img, json=self.simple_txt2img).status_code, 200)
 
+
 class TestTxt2ImgCorrectness(unittest.TestCase):
     pass
 
+
 if __name__ == "__main__":
     unittest.main()