jarviscore-framework 0.1.0__py3-none-any.whl

tests/test_customagent.py

@@ -0,0 +1,248 @@
+"""
+Tests for CustomAgent profile.
+"""
+import pytest
+from jarviscore.profiles.customagent import CustomAgent
+
+
+class ValidCustomAgent(CustomAgent):
+    """Valid CustomAgent with execute_task implementation."""
+    role = "test_custom"
+    capabilities = ["custom_testing"]
+
+    async def execute_task(self, task):
+        return {
+            "status": "success",
+            "output": f"Custom result for: {task.get('task', '')}",
+            "custom_field": "user_data"
+        }
+
+
+class NoExecuteCustomAgent(CustomAgent):
+    """CustomAgent without execute_task (should raise NotImplementedError)."""
+    role = "no_execute"
+    capabilities = ["testing"]
+
+
+class TestCustomAgentInitialization:
+    """Test CustomAgent initialization."""
+
+    def test_valid_customagent_creation(self):
+        """Test creating a valid CustomAgent."""
+        agent = ValidCustomAgent()
+
+        assert agent.role == "test_custom"
+        assert agent.capabilities == ["custom_testing"]
+
+    def test_customagent_with_custom_attributes(self):
+        """Test CustomAgent with user-defined attributes."""
+        class APICustomAgent(CustomAgent):
+            role = "api_agent"
+            capabilities = ["api_calls"]
+            api_endpoint = "https://api.example.com"
+            api_key = "test-key-123"
+
+            async def execute_task(self, task):
+                return {"status": "success"}
+
+        agent = APICustomAgent()
+
+        assert agent.api_endpoint == "https://api.example.com"
+        assert agent.api_key == "test-key-123"
+
+
+class TestCustomAgentSetup:
+    """Test CustomAgent setup."""
+
+    @pytest.mark.asyncio
+    async def test_customagent_setup(self):
+        """Test CustomAgent setup hook."""
+        agent = ValidCustomAgent()
+        await agent.setup()
+
+        # Should run without error
+
+    @pytest.mark.asyncio
+    async def test_customagent_setup_with_framework_initialization(self):
+        """Test CustomAgent setup with framework initialization."""
+        setup_called = []
+
+        class FrameworkCustomAgent(CustomAgent):
+            role = "framework_agent"
+            capabilities = ["framework"]
+
+            async def setup(self):
+                await super().setup()
+                # Simulate initializing a framework (e.g., LangChain, MCP)
+                setup_called.append("framework_initialized")
+
+            async def execute_task(self, task):
+                return {"status": "success"}
+
+        agent = FrameworkCustomAgent()
+        await agent.setup()
+
+        assert "framework_initialized" in setup_called
+
+
+class TestCustomAgentExecution:
+    """Test CustomAgent task execution."""
+
+    @pytest.mark.asyncio
+    async def test_execute_task_implementation(self):
+        """Test CustomAgent execute_task with user implementation."""
+        agent = ValidCustomAgent()
+
+        task = {"task": "Process data"}
+        result = await agent.execute_task(task)
+
+        assert result["status"] == "success"
+        assert "Process data" in result["output"]
+        assert result["custom_field"] == "user_data"
+
+    @pytest.mark.asyncio
+    async def test_execute_task_not_implemented_raises_error(self):
+        """Test that CustomAgent without execute_task raises NotImplementedError."""
+        agent = NoExecuteCustomAgent()
+
+        task = {"task": "Should fail"}
+
+        with pytest.raises(NotImplementedError) as exc_info:
+            await agent.execute_task(task)
+
+        assert "must implement execute_task" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_execute_task_with_optional_cost_tracking(self):
+        """Test CustomAgent returning optional cost tracking fields."""
+        class CostTrackingAgent(CustomAgent):
+            role = "cost_tracker"
+            capabilities = ["tracking"]
+
+            async def execute_task(self, task):
+                return {
+                    "status": "success",
+                    "output": "result",
+                    "tokens_used": 1500,
+                    "cost_usd": 0.003
+                }
+
+        agent = CostTrackingAgent()
+        result = await agent.execute_task({"task": "Track costs"})
+
+        assert result["tokens_used"] == 1500
+        assert result["cost_usd"] == 0.003
+
+
+class TestCustomAgentFrameworkIntegration:
+    """Test CustomAgent with various framework integrations."""
+
+    @pytest.mark.asyncio
+    async def test_langchain_integration_example(self):
+        """Test CustomAgent simulating LangChain integration."""
+        class LangChainAgent(CustomAgent):
+            role = "langchain_agent"
+            capabilities = ["langchain"]
+
+            async def setup(self):
+                await super().setup()
+                # Simulate LangChain initialization
+                self.lc_agent = "mock_langchain_agent"
+
+            async def execute_task(self, task):
+                # Simulate LangChain execution
+                return {
+                    "status": "success",
+                    "output": f"LangChain result: {task['task']}"
+                }
+
+        agent = LangChainAgent()
+        await agent.setup()
+
+        result = await agent.execute_task({"task": "Query database"})
+
+        assert result["status"] == "success"
+        assert "LangChain result" in result["output"]
+
+    @pytest.mark.asyncio
+    async def test_mcp_integration_example(self):
+        """Test CustomAgent simulating MCP integration."""
+        class MCPAgent(CustomAgent):
+            role = "mcp_agent"
+            capabilities = ["mcp_tools"]
+            mcp_server_url = "stdio://./server.py"
+
+            async def setup(self):
+                await super().setup()
+                # Simulate MCP connection
+                self.mcp_client = "mock_mcp_client"
+
+            async def execute_task(self, task):
+                # Simulate MCP tool call
+                return {
+                    "status": "success",
+                    "data": {"tool": "executed", "params": task.get("params")}
+                }
+
+        agent = MCPAgent()
+        await agent.setup()
+
+        result = await agent.execute_task({
+            "task": "Call MCP tool",
+            "params": {"arg1": "value1"}
+        })
+
+        assert result["status"] == "success"
+        assert result["data"]["tool"] == "executed"
+
+    @pytest.mark.asyncio
+    async def test_raw_python_integration_example(self):
+        """Test CustomAgent with raw Python logic."""
+        class DataProcessor(CustomAgent):
+            role = "processor"
+            capabilities = ["data_processing"]
+
+            async def execute_task(self, task):
+                # Pure Python logic
+                data = task.get("params", {}).get("data", [])
+                processed = [x * 2 for x in data]
+                return {
+                    "status": "success",
+                    "output": processed
+                }
+
+        agent = DataProcessor()
+        result = await agent.execute_task({
+            "task": "Double values",
+            "params": {"data": [1, 2, 3, 4, 5]}
+        })
+
+        assert result["status"] == "success"
+        assert result["output"] == [2, 4, 6, 8, 10]
+
+
+class TestCustomAgentInheritance:
+    """Test CustomAgent inheritance from Profile and Agent."""
+
+    def test_customagent_inherits_agent_methods(self):
+        """Test that CustomAgent inherits Agent methods."""
+        agent = ValidCustomAgent()
+
+        # Should have Agent methods
+        assert hasattr(agent, "can_handle")
+        assert hasattr(agent, "execute_task")
+        assert hasattr(agent, "setup")
+        assert hasattr(agent, "teardown")
+
+    def test_customagent_can_handle_tasks(self):
+        """Test that CustomAgent can check task compatibility."""
+        agent = ValidCustomAgent()
+
+        task1 = {"role": "test_custom", "task": "Do something"}
+        assert agent.can_handle(task1) is True
+
+        task2 = {"capability": "custom_testing", "task": "Run tests"}
+        assert agent.can_handle(task2) is True
+
+        task3 = {"role": "different", "task": "Won't handle"}
+        assert agent.can_handle(task3) is False

tests/test_integration.py

@@ -0,0 +1,293 @@
+"""
+Integration tests demonstrating full framework usage.
+"""
+import pytest
+from jarviscore import Mesh, AutoAgent, CustomAgent
+
+
+# Example agents for integration testing
+class ScraperAgent(CustomAgent):
+    """Web scraper agent using CustomAgent profile (mocked for testing)."""
+    role = "scraper"
+    capabilities = ["web_scraping", "data_extraction"]
+
+    async def execute_task(self, task):
+        # Mock scraper for integration testing
+        return {
+            "status": "success",
+            "output": {
+                "url": "example.com",
+                "products": ["Product A", "Product B", "Product C"]
+            },
+            "tokens_used": 0,
+            "cost_usd": 0.0
+        }
+
+
+class ProcessorAgent(CustomAgent):
+    """Data processor using CustomAgent profile."""
+    role = "processor"
+    capabilities = ["data_processing", "transformation"]
+
+    async def execute_task(self, task):
+        # Simulate data processing
+        data = task.get("params", {}).get("data", [])
+        processed = [x.upper() if isinstance(x, str) else x * 2 for x in data]
+
+        return {
+            "status": "success",
+            "output": processed,
+            "tokens_used": 0,  # No LLM usage
+            "cost_usd": 0.0
+        }
+
+
+class StorageAgent(CustomAgent):
+    """Storage agent for saving results."""
+    role = "storage"
+    capabilities = ["database", "file_storage"]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.stored_data = []  # Mock storage
+
+    async def execute_task(self, task):
+        data = task.get("params", {}).get("data")
+        self.stored_data.append(data)
+
+        return {
+            "status": "success",
+            "output": f"Stored {len(self.stored_data)} record(s)",
+            "records_stored": len(self.stored_data)
+        }
+
+
+class TestBasicIntegration:
+    """Test basic framework integration."""
+
+    @pytest.mark.asyncio
+    async def test_single_agent_workflow(self):
+        """Test workflow with single agent."""
+        mesh = Mesh(mode="autonomous")
+        mesh.add(ScraperAgent)
+
+        await mesh.start()
+
+        results = await mesh.workflow("simple-scrape", [
+            {
+                "agent": "scraper",
+                "task": "Scrape example.com for product data"
+            }
+        ])
+
+        assert len(results) == 1
+        assert results[0]["status"] == "success"
+
+        await mesh.stop()
+
+    @pytest.mark.asyncio
+    async def test_multi_agent_workflow(self):
+        """Test workflow with multiple agents."""
+        mesh = Mesh(mode="autonomous")
+        mesh.add(ScraperAgent)
+        mesh.add(ProcessorAgent)
+        mesh.add(StorageAgent)
+
+        await mesh.start()
+
+        results = await mesh.workflow("scrape-process-store", [
+            {
+                "agent": "scraper",
+                "task": "Scrape example.com"
+            },
+            {
+                "agent": "processor",
+                "task": "Process scraped data",
+                "params": {"data": ["hello", "world"]}
+            },
+            {
+                "agent": "storage",
+                "task": "Save processed data",
+                "params": {"data": ["HELLO", "WORLD"]}
+            }
+        ])
+
+        assert len(results) == 3
+        assert all(r["status"] == "success" for r in results)
+
+        await mesh.stop()
+
+
+class TestMixedProfileIntegration:
+    """Test integration of AutoAgent and CustomAgent profiles."""
+
+    @pytest.mark.asyncio
+    async def test_auto_and_custom_agents_together(self):
+        """Test that AutoAgent and CustomAgent work together."""
+        mesh = Mesh(mode="autonomous")
+
+        # Add AutoAgent
+        scraper = mesh.add(ScraperAgent)
+
+        # Add CustomAgent
+        processor = mesh.add(ProcessorAgent)
+
+        await mesh.start()
+
+        # Verify both agents are registered
+        assert mesh.get_agent("scraper") == scraper
+        assert mesh.get_agent("processor") == processor
+
+        # Execute workflow using both
+        results = await mesh.workflow("mixed-workflow", [
+            {"agent": "scraper", "task": "Scrape data"},
+            {"agent": "processor", "task": "Process", "params": {"data": [1, 2, 3]}}
+        ])
+
+        assert len(results) == 2
+        assert results[0]["status"] == "success"
+        assert results[1]["status"] == "success"
+
+        await mesh.stop()
+
+
+class TestCapabilityRouting:
+    """Test task routing by capabilities."""
+
+    @pytest.mark.asyncio
+    async def test_route_by_capability(self):
+        """Test that tasks are routed by capability."""
+        mesh = Mesh(mode="autonomous")
+        mesh.add(ProcessorAgent)
+
+        await mesh.start()
+
+        # Route by capability instead of role
+        results = await mesh.workflow("capability-routing", [
+            {
+                "agent": "data_processing",  # Capability, not role
+                "task": "Process this data",
+                "params": {"data": [10, 20, 30]}
+            }
+        ])
+
+        assert len(results) == 1
+        assert results[0]["status"] == "success"
+        assert results[0]["output"] == [20, 40, 60]  # Doubled
+
+        await mesh.stop()
+
+
+class TestAgentStateAndCustomization:
+    """Test agent state management and customization."""
+
+    @pytest.mark.asyncio
+    async def test_agent_maintains_state_across_tasks(self):
+        """Test that CustomAgent can maintain state across multiple tasks."""
+        mesh = Mesh(mode="autonomous")
+        storage = mesh.add(StorageAgent)
+
+        await mesh.start()
+
+        # Execute multiple storage tasks
+        await mesh.workflow("multi-store", [
+            {"agent": "storage", "task": "Store item 1", "params": {"data": "item1"}},
+            {"agent": "storage", "task": "Store item 2", "params": {"data": "item2"}},
+            {"agent": "storage", "task": "Store item 3", "params": {"data": "item3"}},
+        ])
+
+        # Verify storage agent maintained state
+        assert len(storage.stored_data) == 3
+        assert storage.stored_data == ["item1", "item2", "item3"]
+
+        await mesh.stop()
+
+
+class TestErrorHandling:
+    """Test error handling in workflows."""
+
+    @pytest.mark.asyncio
+    async def test_missing_agent_error(self):
+        """Test workflow with missing agent."""
+        mesh = Mesh(mode="autonomous")
+        mesh.add(ProcessorAgent)
+
+        await mesh.start()
+
+        results = await mesh.workflow("missing-agent", [
+            {
+                "agent": "nonexistent_agent",
+                "task": "Should fail"
+            }
+        ])
+
+        assert len(results) == 1
+        assert results[0]["status"] == "failure"
+        assert "No agent found" in results[0]["error"]
+
+        await mesh.stop()
+
+
+class TestRealWorldScenario:
+    """Test real-world usage scenarios."""
+
+    @pytest.mark.asyncio
+    async def test_data_pipeline_scenario(self):
+        """Test complete data pipeline: scrape → process → store."""
+        mesh = Mesh(mode="autonomous")
+        mesh.add(ScraperAgent)
+        mesh.add(ProcessorAgent)
+        storage = mesh.add(StorageAgent)
+
+        await mesh.start()
+
+        # Simulate data pipeline
+        results = await mesh.workflow("data-pipeline", [
+            {
+                "agent": "scraper",
+                "task": "Scrape example.com for products"
+            },
+            {
+                "agent": "processor",
+                "task": "Clean and normalize product data",
+                "params": {"data": ["product-a", "product-b", "product-c"]},
+                "depends_on": [0]
+            },
+            {
+                "agent": "storage",
+                "task": "Save to database",
+                "params": {"data": ["PRODUCT-A", "PRODUCT-B", "PRODUCT-C"]},
+                "depends_on": [1]
+            }
+        ])
+
+        # Verify all steps succeeded
+        assert len(results) == 3
+        assert all(r["status"] == "success" for r in results)
+
+        # Verify data was stored
+        assert len(storage.stored_data) == 1
+        assert storage.stored_data[0] == ["PRODUCT-A", "PRODUCT-B", "PRODUCT-C"]
+
+        await mesh.stop()
+
+    @pytest.mark.asyncio
+    async def test_agent_discovery_by_capability(self):
+        """Test that mesh can discover agents by capability."""
+        mesh = Mesh(mode="autonomous")
+        mesh.add(ScraperAgent)
+        mesh.add(ProcessorAgent)
+        mesh.add(StorageAgent)
+
+        # Test capability indexing
+        scrapers = mesh.get_agents_by_capability("web_scraping")
+        assert len(scrapers) == 1
+        assert scrapers[0].role == "scraper"
+
+        processors = mesh.get_agents_by_capability("data_processing")
+        assert len(processors) == 1
+        assert processors[0].role == "processor"
+
+        storage_agents = mesh.get_agents_by_capability("database")
+        assert len(storage_agents) == 1
+        assert storage_agents[0].role == "storage"