agents/testing_agent.py

from typing import Dict, Any, List
from utils.llm_orchestrator import LLMOrchestrator
from loguru import logger
from utils import llm_orchestrator
import unittest
import io
import contextlib
import asyncio
import sys

class TestingAgent:
    def __init__(self, llm_orchestrator: LLMOrchestrator):
        """Initialize the Testing Agent."""
        logger.info("Initializing TestingAgent")
        self.llm_orchestrator = llm_orchestrator
        self.capabilities = [
            "test_generation",
            "test_execution",
            "test_analysis"
        ]
        self.setup_logger()

    def setup_logger(self):
        """Configure logging for the agent."""
        logger.add("logs/testing_agent.log", rotation="500 MB")

    async def generate_tests(
            self, code: str, focus_areas: List[str] = None) -> Dict[str, Any]:
        """Generate test cases for the given code."""
        logger.info(f"Generating tests for code:\n{code}")
        try:
            prompt = f"""
            Generate test cases for the following code:
            ```python
            {code}
            ```

            Focus on the following areas:
            {', '.join(focus_areas) if focus_areas else 'General functionality'}

            Provide the test cases in a format compatible with Python's unittest framework.
            """
            test_code = await self.llm_orchestrator.generate_completion(prompt)

            logger.info(f"Test cases generated:\n{test_code}")
            return {
                "status": "success",
                "test_code": test_code
            }

        except Exception as e:
            logger.error(f"Error generating tests: {str(e)}")
            return {
                "status": "error",
                "message": str(e)
            }

    async def execute_tests(self, test_code: str) -> Dict[str, Any]:
        """Execute the generated test cases in a separate process."""
        logger.info(f"Executing tests in a separate process:\n{test_code}")
        try:
            process = await asyncio.create_subprocess_exec(
                sys.executable,  # Use the same Python interpreter
                "agents/test_runner.py",
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )

            stdout, stderr = await process.communicate(input=test_code.encode())
            test_output = stdout.decode()
            error_output = stderr.decode()

            if process.returncode == 0:
                logger.info(f"Test execution completed:\n{test_output}")
                return {
                    "status": "success",
                    "results": test_output
                }
            else:
                logger.error(f"Error executing tests:\n{error_output}")
                return {
                    "status": "error",
                    "message": error_output
                }
        except Exception as e:
            logger.error(f"Error executing tests: {str(e)}")
            return {
                "status": "error",
                "message": str(e)
            }