From 28a9df67652428c87e39b5ffaca4868a35142c0d Mon Sep 17 00:00:00 2001 From: Andres Caicedo Date: Thu, 20 Apr 2023 20:58:42 +0200 Subject: [PATCH] Update test_image_gen.py --- tests/test_image_gen.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_image_gen.py b/tests/test_image_gen.py index 19c57e427d5c..dd261739234b 100644 --- a/tests/test_image_gen.py +++ b/tests/test_image_gen.py @@ -10,15 +10,19 @@ def lst(txt): + """Get the last item in a list.""" return txt.split(":")[1].strip() @unittest.skipIf(os.getenv("CI"), "Skipping image generation tests") class TestImageGen(unittest.TestCase): + """Test image generation.""" def setUp(self): + """Set up the test class.""" self.config = Config() def test_dalle(self): + """Test DALL-E image generation.""" self.config.image_provider = "dalle" # Test using size 256 @@ -37,6 +41,7 @@ def test_dalle(self): image_path.unlink() def test_huggingface(self): + """Test HuggingFace image generation.""" self.config.image_provider = "huggingface" # Test usin SD 1.4 model and size 512 @@ -57,6 +62,7 @@ def test_huggingface(self): image_path.unlink() def test_sd_webui(self): + """Test SD WebUI image generation.""" self.config.image_provider = "sd_webui" return @@ -99,4 +105,5 @@ def test_sd_webui(self): if __name__ == "__main__": + """Run the tests.""" unittest.main()