Added Handler - execute command - test

src/main.py

@@ -1,4 +1,5 @@
 from src import server
 
-server = server.Server("localhost", 12345)
-server.start()
+if __name__ == "__main__":
+    server = server.Server("localhost", 12345)
+    server.start()

src/utils.py

@@ -90,19 +90,23 @@ def build_conditional_gan(generator, discriminator):
     Returns:
         conditionalGAN: Compiled cGAN model.
     """
-    cond_gan = cgan.ConditionalGAN(
+    config = cgan.GANConfig(
         discriminator=discriminator,
         generator=generator,
         latent_dim=latent_dim,
         image_size=image_size,
         num_classes=num_classes,
     )
-    cond_gan.compile(
+
+    cond_gan = cgan.ConditionalGAN(
+        config=config,
         d_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
         g_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
         loss_fn=keras.losses.BinaryCrossentropy(from_logits=True),
     )
 
+    cond_gan.compile()
+
     return cond_gan
 
 

tests/test_handler.py

@@ -0,0 +1,65 @@
+import os
+import unittest
+from unittest.mock import MagicMock
+
+from src.handler import Handler
+
+
+class TestHandler(unittest.TestCase):
+    def test_execute_command_generate_number_success(self):
+        # Arrange
+        request = {
+            "command": "generate_number",
+            "text": "5"
+        }
+        expected_file_path = "./data/nums/drawn_number_5.png"
+        handler = Handler(MagicMock())
+
+        # Act
+        handler._execute_command(request)
+
+        # Assert
+        self.assertTrue(os.path.exists(expected_file_path))
+
+        # Clean up
+        os.remove(expected_file_path)
+
+
+    def test_execute_command_generate_number_error(self):
+        # Arrange
+        request = {
+            "command": "generate_number",
+            "text": "abc"
+        }
+        expected_response = {
+            "status": "error",
+            "message": "Error generating the image: invalid literal for int() with base 10: 'abc'"
+        }
+        handler = Handler(MagicMock())
+
+        # Act
+        response = handler._execute_command(request)
+
+        # Assert
+        self.assertEqual(response, expected_response)
+
+    def test_execute_command_unknown_command(self):
+        # Arrange
+        request = {
+            "command": "unknown_command",
+            "text": "5"
+        }
+        expected_response = {
+            "status": "error",
+            "message": "Unknown command: unknown_command"
+        }
+        handler = Handler(MagicMock())
+
+        # Act
+        response = handler._execute_command(request)
+
+        # Assert
+        self.assertEqual(response, expected_response)
+
+if __name__ == '__main__':
+    unittest.main()