diff --git a/smolswarms/swarm_factory.py b/smolswarms/swarm_factory.py
index 6172cf8..b088693 100644
--- a/smolswarms/swarm_factory.py
+++ b/smolswarms/swarm_factory.py
@@ -1,17 +1,20 @@
-from typing import Dict, Any
+from typing import Any
+from typing import Dict
+from typing import Optional
+
 
 class SwarmFactory:
     """
     Factory class for spawning different types of agent swarms.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """
         Initializes the SwarmFactory.
         """
         pass
 
-    def spawn_department(self, config: Dict[str, Any]):
+    def spawn_department(self, config: Dict[str, Any]) -> Optional[Any]:
         """
         Spawns a department of agents based on the given configuration.
 
@@ -19,7 +22,7 @@ def spawn_department(self, config: Dict[str, Any]):
             config (Dict[str, Any]): Configuration for the department.
 
         Returns:
-            Any: The spawned department.
+            Optional[Any]: The spawned department or None if creation fails.
         """
         business_unit = config.get("business_unit")
         budget = config.get("budget")
@@ -32,13 +35,19 @@ def spawn_department(self, config: Dict[str, Any]):
             model = "gpt-4"
 
         print(f"Selected model: {model}")
-
-        print(f"Spawning department: {business_unit} with budget: {budget} and vibe: {vibe}")
+        print(
+            f"Spawning department: {business_unit} with budget: {budget} \
+              and vibe: {vibe}"
+        )
         return None
+
+
 if __name__ == "__main__":
     factory = SwarmFactory()
-    swarm = factory.spawn_department({
-        "business_unit": "growth_team",
-        "budget": "optimize_tokens",
-        "vibe": "hypergrowth"
-    })
+    swarm = factory.spawn_department(
+        {
+            "business_unit": "growth_team",
+            "budget": "optimize_tokens",
+            "vibe": "hypergrowth",
+        }
+    )
diff --git a/tests/test_swarm_factory.py b/tests/test_swarm_factory.py
index 640cc6d..3f8a4b6 100644
--- a/tests/test_swarm_factory.py
+++ b/tests/test_swarm_factory.py
@@ -1,24 +1,29 @@
-import unittest
+"""Tests for SwarmFactory."""
+
 from smolswarms.swarm_factory import SwarmFactory
 
-class TestSwarmFactory(unittest.TestCase):
-    def test_spawn_department_optimize_tokens(self):
-        factory = SwarmFactory()
-        swarm = factory.spawn_department({
-            "business_unit": "test_team",
+
+def test_spawn_department() -> None:
+    """Test spawning a department."""
+    factory = SwarmFactory()
+    swarm = factory.spawn_department(
+        {
+            "business_unit": "growth_team",
             "budget": "optimize_tokens",
-            "vibe": "test_vibe"
-        })
-        self.assertIsNone(swarm)  # Add more specific assertions based on expected behavior
+            "vibe": "hypergrowth",
+        }
+    )
+    assert swarm is None
 
-    def test_spawn_department_default_budget(self):
-        factory = SwarmFactory()
-        swarm = factory.spawn_department({
-            "business_unit": "test_team",
-            "budget": "default",
-            "vibe": "test_vibe"
-        })
-        self.assertIsNone(swarm)  # Add more specific assertions based on expected behavior
 
-if __name__ == '__main__':
-    unittest.main()
+def test_spawn_department_with_different_budget() -> None:
+    """Test spawning a department with a different budget."""
+    factory = SwarmFactory()
+    swarm = factory.spawn_department(
+        {
+            "business_unit": "growth_team",
+            "budget": "unlimited",
+            "vibe": "hypergrowth",
+        }
+    )
+    assert swarm is None