agent_lab/backend/test_ech0_stack.py

import unittest
import os
import json
import shutil
from unittest.mock import patch, MagicMock
import sys

# Import modules to test
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))

from agent_lab.backend.universal_lab import universal_lab
from agent_lab.backend.lab_notebook import lab_notebook
from agent_lab.backend.ech0_service import ech0
from agent_lab.backend.main import app

try:
    from fastapi.testclient import TestClient
    HAS_TEST_CLIENT = True
except (ImportError, RuntimeError):
    HAS_TEST_CLIENT = False
    print("WARNING: 'httpx' or 'fastapi' not found. API Endpoint tests will be skipped.")

class TestECH0Stack(unittest.TestCase):
    def setUp(self):
        self.original_log_dir = lab_notebook._ensure_log_dir
        if HAS_TEST_CLIENT:
            self.client = TestClient(app)

    def test_universal_lab_discovery(self):
        """Test if Universal Lab correctly finds lab files."""
        print("\n[TEST] Universal Lab Discovery")
        labs = universal_lab.labs
        self.assertTrue(len(labs) > 0)
        print(f"Verified {len(labs)} labs discovered.")

    def test_material_combination_and_nist(self):
        """Test Na+Cl logic and NIST validation."""
        print("\n[TEST] Universal Lab Combination & NIST")
        
        # Test 1: Sodium + Chlorine (NIST Verified)
        res = universal_lab.combine_materials(["Sodium", "Chlorine"], {"temperature": 25})
        self.assertTrue(res["reaction_occurred"])
        self.assertIn("Sodium Chloride", res["products"])
        self.assertTrue(res["validation"]["verified"])
        self.assertIn("NIST Ref", res["validation"]["source"])
        print("Verified Na+Cl reaction and NIST check.")

        # Test 2: Generic fallback
        res = universal_lab.combine_materials(["Iron", "Wood"], {"temperature": 25})
        self.assertFalse(res["reaction_occurred"])
        self.assertFalse(res["validation"]["verified"])
        print("Verified generic fallback.")

    def test_lab_notebook_persistence(self):
        """Test that experiments are logged to notebook."""
        print("\n[TEST] Lab Notebook Persistence")
        
        # Log a dummy experiment
        data = {"test_id": 12345, "timestamp": "now", "data": "test"}
        lab_notebook.log_experiment(data)
        
        # Fetch logs
        logs = lab_notebook.get_logs()
        # Check if our data is in there (the last one might not be ours if parallel, but likely is)
        found = False
        for entry in logs[-5:]:
            if entry.get("data", {}).get("test_id") == 12345:
                found = True
                break
        
        self.assertTrue(found, "Notebook entry not found.")
        print("Verified notebook entry creation.")

    def test_api_endpoints(self):
        """Test the FastAPI endpoints."""
        if not HAS_TEST_CLIENT:
            print("\n[TEST] API Endpoints: SKIPPED (Missing dependencies)")
            return

        print("\n[TEST] API Endpoints")
        
        # Test Greeting
        with patch.object(ech0, 'generate_greeting', return_value="Test Greeting"):
            response = self.client.get("/chat/greeting")
            self.assertEqual(response.status_code, 200)
            self.assertEqual(response.json(), {"greeting": "Test Greeting"})
            
        # Test Notebook endpoint
        response = self.client.get("/notebook")
        self.assertEqual(response.status_code, 200)
        self.assertIsInstance(response.json(), list)
        print("Verified /chat/greeting and /notebook endpoints.")

    def test_ech0_logic_flow(self):
        """Test ECH0 tool parsing logic."""
        print("\n[TEST] ECH0 Tool Parsing")
        
        # Mock the LLM response to simulate a JSON tool call
        mock_json_response = """
        Thinking about it...
        ```json
        {
            "tool": "combine_materials",
            "materials": ["Hydrogen", "Oxygen"],
            "conditions": {"temperature": 100}
        }
        ```
        """
        
        with patch.object(ech0, '_call_llm', return_value=mock_json_response):
            result = ech0.chat("Make water")
            self.assertEqual(result["tool_used"], "combine_materials")
            self.assertIsNotNone(result["action_result"])
            print("Verified ECH0 JSON parsing.")

if __name__ == '__main__':
    unittest.main()