Update test_image_gen.py

yoitzme · Apr 20, 2023 · 28a9df6 · 28a9df6
1 parent b169455
commit 28a9df6
Showing 1 changed file with 7 additions and 0 deletions.
diff --git 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()