Add unit testing

Author: GustyCube

easilyai/enhanced_app.py

@@ -7,7 +7,7 @@
 """
 
 import time
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 from .app import EasyAIApp, create_app as _create_app
 from .config import get_config, EasilyAIConfig
 from .cache import get_cache

tests/test_anthropic_service.py

@@ -1,28 +1,66 @@
-# Disabled Temporarily due to the curent nature of the code. 
+import unittest
+from unittest.mock import patch, MagicMock
+from easilyai.services.anthropic_service import AnthropicService
+from easilyai.exceptions import MissingAPIKeyError, AuthenticationError, ServerError
 
-# from unittest import TestCase
-# from unittest.mock import Mock, patch
-# from easilyai.services.anthropic_service import AnthropicService
-# import anthropic
 
+class TestAnthropicService(unittest.TestCase):
+    def setUp(self):
+        self.apikey = "fake_api_key"
+        self.model = "claude-3-sonnet-20240229"
+        
+    def test_missing_api_key(self):
+        with self.assertRaises(MissingAPIKeyError):
+            AnthropicService(apikey=None, model=self.model)
 
-# class TestAnthropicService(TestCase):
-#     def setUp(self):
-#         self.service = AnthropicService(apikey="test_api_key", model="claude-3-5", max_tokens=1024)
+    @patch.object(AnthropicService, '__init__', lambda x, y, z, **kwargs: None)
+    @patch('anthropic.Anthropic')
+    def test_generate_text_success(self, mock_anthropic_class):
+        mock_client = mock_anthropic_class.return_value
+        mock_response = MagicMock()
+        mock_response.content = [MagicMock(text="Mocked Anthropic response")]
+        mock_client.messages.create.return_value = mock_response
+        
+        service = AnthropicService("fake_key", "claude-3-sonnet-20240229")
+        service.client = mock_client
+        service.model = "claude-3-sonnet-20240229"
+        service.max_tokens = 1024
+        
+        response = service.generate_text("Test prompt")
+        self.assertEqual(response, "Mocked Anthropic response")
 
-#     @patch('anthropic.Anthropic.messages.create', new_callable=Mock)
-#     def test_generate_text(self, mock_messages):
-#         mock_messages.create.return_value = {
-#             "content": [{"text": "Mocked response"}]
-#         }
+    @patch.object(AnthropicService, '__init__', lambda x, y, z, **kwargs: None)
+    @patch('anthropic.Anthropic')
+    def test_generate_text_with_image(self, mock_anthropic_class):
+        mock_client = mock_anthropic_class.return_value
+        mock_response = MagicMock()
+        mock_response.content = [MagicMock(text="Mocked response with image")]
+        mock_client.messages.create.return_value = mock_response
+        
+        service = AnthropicService("fake_key", "claude-3-sonnet-20240229")
+        service.client = mock_client
+        service.model = "claude-3-sonnet-20240229"
+        service.max_tokens = 1024
+        
+        # Mock prepare_image to return None (simulating URL instead of local file)
+        with patch.object(service, 'prepare_image', return_value=None):
+            response = service.generate_text("Describe this image", "http://example.com/image.jpg")
+            self.assertEqual(response, "Mocked response with image")
 
-#         response = self.service.generate_text("Test prompt")
+    @patch.object(AnthropicService, '__init__', lambda x, y, z, **kwargs: None)
+    @patch('anthropic.Anthropic')
+    def test_generate_text_authentication_error(self, mock_anthropic_class):
+        mock_client = mock_anthropic_class.return_value
+        mock_client.messages.create.side_effect = Exception("Authentication failed")
+        
+        service = AnthropicService("fake_key", "claude-3-sonnet-20240229")
+        service.client = mock_client
+        service.model = "claude-3-sonnet-20240229"
+        service.max_tokens = 1024
+        
+        with self.assertRaises(ServerError):
+            service.generate_text("Test prompt")
 
-#         self.assertEqual(response, "Mocked response")
 
-#     @patch('anthropic.Anthropic.messages.create', new_callable=Mock)
-#     def test_generate_text_error(self, mock_messages):
-#         mock_messages.create.side_effect = Exception("API Error")
-
-#         with self.assertRaises(Exception):
-#             self.service.generate_text("Test prompt")
+if __name__ == "__main__":
+    unittest.main()

tests/test_batch.py

@@ -0,0 +1,31 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+
+class TestBatchProcessing(unittest.TestCase):
+    
+    def test_batch_module_imports(self):
+        # Test that batch module can be imported
+        try:
+            from easilyai import batch
+            self.assertIsNotNone(batch)
+        except ImportError:
+            self.skipTest("Batch module not available")
+
+    def test_batch_functionality_exists(self):
+        # Test that batch processing functionality exists
+        try:
+            from easilyai.batch import BatchProcessor
+            self.assertTrue(hasattr(BatchProcessor, '__init__'))
+        except ImportError:
+            # If BatchProcessor doesn't exist, check for other batch functions
+            try:
+                import easilyai.batch
+                # At least the module should exist
+                self.assertIsNotNone(easilyai.batch)
+            except ImportError:
+                self.skipTest("No batch processing functionality found")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_custom_ai.py

@@ -0,0 +1,86 @@
+import unittest
+from easilyai.custom_ai import CustomAIService, register_custom_ai, _registered_custom_ais
+
+
+class TestCustomAI(unittest.TestCase):
+    def setUp(self):
+        # Clear registered custom AIs before each test
+        _registered_custom_ais.clear()
+
+    def test_custom_ai_service_init(self):
+        service = CustomAIService(model="test-model", apikey="test-key")
+        self.assertEqual(service.model, "test-model")
+        self.assertEqual(service.apikey, "test-key")
+
+    def test_custom_ai_service_init_without_apikey(self):
+        service = CustomAIService(model="test-model")
+        self.assertEqual(service.model, "test-model")
+        self.assertIsNone(service.apikey)
+
+    def test_custom_ai_service_not_implemented_methods(self):
+        service = CustomAIService(model="test-model")
+        
+        with self.assertRaises(NotImplementedError):
+            service.generate_text("test prompt")
+        
+        with self.assertRaises(NotImplementedError):
+            service.generate_image("test prompt")
+        
+        with self.assertRaises(NotImplementedError):
+            service.text_to_speech("test text")
+
+    def test_register_valid_custom_ai(self):
+        class ValidCustomAI(CustomAIService):
+            def generate_text(self, prompt):
+                return f"Generated: {prompt}"
+        
+        register_custom_ai("valid_ai", ValidCustomAI)
+        self.assertIn("valid_ai", _registered_custom_ais)
+        self.assertEqual(_registered_custom_ais["valid_ai"], ValidCustomAI)
+
+    def test_register_invalid_custom_ai(self):
+        class InvalidCustomAI:
+            pass
+        
+        with self.assertRaises(TypeError) as context:
+            register_custom_ai("invalid_ai", InvalidCustomAI)
+        
+        self.assertIn("Custom service must inherit from CustomAIService", str(context.exception))
+        self.assertNotIn("invalid_ai", _registered_custom_ais)
+
+    def test_register_multiple_custom_ais(self):
+        class CustomAI1(CustomAIService):
+            def generate_text(self, prompt):
+                return "AI1 response"
+        
+        class CustomAI2(CustomAIService):
+            def generate_text(self, prompt):
+                return "AI2 response"
+        
+        register_custom_ai("ai1", CustomAI1)
+        register_custom_ai("ai2", CustomAI2)
+        
+        self.assertEqual(len(_registered_custom_ais), 2)
+        self.assertIn("ai1", _registered_custom_ais)
+        self.assertIn("ai2", _registered_custom_ais)
+
+    def test_custom_ai_implementation_example(self):
+        class MockCustomAI(CustomAIService):
+            def generate_text(self, prompt):
+                return f"Mock response to: {prompt}"
+            
+            def generate_image(self, prompt):
+                return f"Mock image for: {prompt}"
+            
+            def text_to_speech(self, text):
+                return f"Mock audio for: {text}"
+        
+        service = MockCustomAI(model="mock-model", apikey="mock-key")
+        
+        self.assertEqual(service.generate_text("Hello"), "Mock response to: Hello")
+        self.assertEqual(service.generate_image("Cat"), "Mock image for: Cat")
+        self.assertEqual(service.text_to_speech("Speech"), "Mock audio for: Speech")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_enhanced_app.py

@@ -0,0 +1,31 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+
+class TestEnhancedApp(unittest.TestCase):
+
+    def test_enhanced_app_imports(self):
+        # Test that enhanced app modules can be imported
+        try:
+            from easilyai.enhanced_app import create_enhanced_app
+            self.assertTrue(callable(create_enhanced_app))
+        except ImportError:
+            self.fail("Could not import enhanced_app module")
+
+    def test_enhanced_app_function_signature(self):
+        # Test the function exists with correct parameters
+        from easilyai.enhanced_app import create_enhanced_app
+        import inspect
+        
+        sig = inspect.signature(create_enhanced_app)
+        params = list(sig.parameters.keys())
+        
+        # Check required parameters exist
+        self.assertIn('name', params)
+        self.assertIn('service', params)
+        self.assertIn('api_key', params)
+        self.assertIn('model', params)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_enhanced_pipeline.py

@@ -0,0 +1,129 @@
+import unittest
+from unittest.mock import MagicMock, patch
+from easilyai.enhanced_pipeline import (
+    EnhancedPipeline, TaskStatus, ExecutionMode, TaskResult, PipelineTask
+)
+
+
+class TestEnhancedPipeline(unittest.TestCase):
+    def setUp(self):
+        self.mock_app = MagicMock()
+        self.pipeline = EnhancedPipeline("test_pipeline")
+
+    def test_pipeline_init(self):
+        self.assertEqual(self.pipeline.name, "test_pipeline")
+        self.assertEqual(len(self.pipeline.tasks), 0)
+        self.assertEqual(self.pipeline.execution_mode, ExecutionMode.SEQUENTIAL)
+
+    def test_add_task_simple(self):
+        task_id = self.pipeline.add_task("task1", self.mock_app, "generate_text", "Hello")
+        self.assertEqual(len(self.pipeline.tasks), 1)
+        self.assertIsInstance(task_id, str)
+        
+        task = self.pipeline.tasks[task_id]
+        self.assertEqual(task.task_type, "generate_text")
+        self.assertEqual(task.prompt, "Hello")
+        self.assertEqual(task.status, TaskStatus.PENDING)
+
+    def test_add_task_with_dependencies(self):
+        task1_id = self.pipeline.add_task("task1", self.mock_app, "generate_text", "Task 1")
+        task2_id = self.pipeline.add_task("task2", self.mock_app, "generate_text", "Task 2", dependencies=[task1_id])
+        
+        self.assertEqual(len(self.pipeline.tasks), 2)
+        task2 = self.pipeline.tasks[task2_id]
+        self.assertEqual(task2.dependencies, [task1_id])
+
+    def test_add_task_with_condition(self):
+        condition = lambda results: True
+        task_id = self.pipeline.add_task("task1", self.mock_app, "generate_text", "Conditional task", condition=condition)
+        
+        task = self.pipeline.tasks[task_id]
+        self.assertEqual(task.condition, condition)
+
+    def test_task_status_enum(self):
+        self.assertEqual(TaskStatus.PENDING.value, "pending")
+        self.assertEqual(TaskStatus.RUNNING.value, "running")
+        self.assertEqual(TaskStatus.COMPLETED.value, "completed")
+        self.assertEqual(TaskStatus.FAILED.value, "failed")
+        self.assertEqual(TaskStatus.SKIPPED.value, "skipped")
+
+    def test_execution_mode_enum(self):
+        self.assertEqual(ExecutionMode.SEQUENTIAL.value, "sequential")
+        self.assertEqual(ExecutionMode.PARALLEL.value, "parallel")
+        self.assertEqual(ExecutionMode.CONDITIONAL.value, "conditional")
+
+    def test_task_result_creation(self):
+        result = TaskResult(
+            task_id="test-task",
+            status=TaskStatus.COMPLETED,
+            result="Test result",
+            duration=1.5
+        )
+        
+        self.assertEqual(result.task_id, "test-task")
+        self.assertEqual(result.status, TaskStatus.COMPLETED)
+        self.assertEqual(result.result, "Test result")
+        self.assertEqual(result.duration, 1.5)
+        self.assertIsNone(result.error)
+        self.assertEqual(result.metadata, {})
+
+    def test_set_execution_mode(self):
+        self.pipeline.set_execution_mode(ExecutionMode.PARALLEL)
+        self.assertEqual(self.pipeline.execution_mode, ExecutionMode.PARALLEL)
+
+    def test_clear_tasks(self):
+        self.pipeline.add_task("task1", self.mock_app, "generate_text", "Task 1")
+        self.pipeline.add_task("task2", self.mock_app, "generate_text", "Task 2")
+        self.assertEqual(len(self.pipeline.tasks), 2)
+        
+        self.pipeline.tasks.clear()  # Simplified clear method
+        self.assertEqual(len(self.pipeline.tasks), 0)
+
+    @patch('easilyai.enhanced_pipeline.time.time')
+    def test_simple_execution_simulation(self, mock_time):
+        # Mock time for duration calculation
+        mock_time.side_effect = [0.0, 1.0]  # Start and end times
+        
+        # Mock app response
+        self.mock_app.request.return_value = "Generated response"
+        
+        # Add a simple task
+        task_id = self.pipeline.add_task("task1", self.mock_app, "generate_text", "Hello")
+        
+        # This test verifies the structure, but actual execution would require
+        # the full implementation which might involve async/threading
+        self.assertEqual(len(self.pipeline.tasks), 1)
+        task = self.pipeline.tasks[task_id]
+        self.assertEqual(task.task_id, task_id)
+        self.assertEqual(task.status, TaskStatus.PENDING)
+
+    def test_variable_substitution_pattern(self):
+        # Test that the pipeline can handle variable patterns
+        prompt_with_vars = "Hello {name}, how are you?"
+        task_id = self.pipeline.add_task("task1", self.mock_app, "generate_text", prompt_with_vars)
+        
+        task = self.pipeline.tasks[task_id]
+        self.assertIn("{name}", task.prompt)
+
+    def test_pipeline_task_attributes(self):
+        task_id = self.pipeline.add_task(
+            "task1",
+            self.mock_app,
+            "generate_text", 
+            "Test", 
+            dependencies=["dep1"], 
+            condition=lambda x: True,
+            retry_count=3,
+            timeout=30
+        )
+        
+        task = self.pipeline.tasks[task_id]
+        self.assertEqual(task.task_type, "generate_text")
+        self.assertEqual(task.dependencies, ["dep1"])
+        self.assertIsNotNone(task.condition)
+        self.assertEqual(task.retry_count, 3)
+        self.assertEqual(task.timeout, 30)
+
+
+if __name__ == "__main__":
+    unittest.main()