123 lines
5.2 KiB
Python
123 lines
5.2 KiB
Python
import unittest
|
|
from fastapi.testclient import TestClient
|
|
import gui_2
|
|
from unittest.mock import patch, MagicMock
|
|
from pathlib import Path
|
|
|
|
class TestHeadlessAPI(unittest.TestCase):
|
|
def setUp(self):
|
|
# We need an App instance to initialize the API, but we want to avoid GUI stuff
|
|
with patch('gui_2.session_logger.open_session'), \
|
|
patch('gui_2.ai_client.set_provider'), \
|
|
patch('gui_2.session_logger.close_session'):
|
|
self.app_instance = gui_2.App()
|
|
# Set a default API key for tests
|
|
self.test_api_key = "test-secret-key"
|
|
self.app_instance.config["headless"] = {"api_key": self.test_api_key}
|
|
self.headers = {"X-API-KEY": self.test_api_key}
|
|
|
|
# Clear any leftover state
|
|
self.app_instance._pending_actions = {}
|
|
self.app_instance._pending_dialog = None
|
|
|
|
self.api = self.app_instance.create_api()
|
|
self.client = TestClient(self.api)
|
|
|
|
def test_health_endpoint(self):
|
|
response = self.client.get("/health")
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertEqual(response.json(), {"status": "ok"})
|
|
|
|
def test_status_endpoint_unauthorized(self):
|
|
# Ensure a key is required
|
|
with patch.dict(self.app_instance.config, {"headless": {"api_key": "some-required-key"}}):
|
|
response = self.client.get("/status")
|
|
self.assertEqual(response.status_code, 403)
|
|
|
|
def test_status_endpoint_authorized(self):
|
|
# We'll use a test key
|
|
headers = {"X-API-KEY": "test-secret-key"}
|
|
with patch.dict(self.app_instance.config, {"headless": {"api_key": "test-secret-key"}}):
|
|
response = self.client.get("/status", headers=headers)
|
|
self.assertEqual(response.status_code, 200)
|
|
|
|
def test_generate_endpoint(self):
|
|
payload = {
|
|
"prompt": "Hello AI"
|
|
}
|
|
# Mock ai_client.send and get_comms_log
|
|
with patch('gui_2.ai_client.send') as mock_send, \
|
|
patch('gui_2.ai_client.get_comms_log') as mock_log:
|
|
mock_send.return_value = "Hello from Mock AI"
|
|
mock_log.return_value = [{
|
|
"kind": "response",
|
|
"payload": {
|
|
"usage": {"input_tokens": 10, "output_tokens": 5}
|
|
}
|
|
}]
|
|
|
|
response = self.client.post("/api/v1/generate", json=payload, headers=self.headers)
|
|
self.assertEqual(response.status_code, 200)
|
|
data = response.json()
|
|
self.assertEqual(data["text"], "Hello from Mock AI")
|
|
self.assertIn("metadata", data)
|
|
self.assertEqual(data["usage"]["input_tokens"], 10)
|
|
|
|
def test_pending_actions_endpoint(self):
|
|
# Manually add a pending action
|
|
with patch('gui_2.uuid.uuid4', return_value="test-action-id"):
|
|
dialog = gui_2.ConfirmDialog("dir", ".")
|
|
self.app_instance._pending_actions[dialog._uid] = dialog
|
|
|
|
response = self.client.get("/api/v1/pending_actions", headers=self.headers)
|
|
self.assertEqual(response.status_code, 200)
|
|
data = response.json()
|
|
self.assertEqual(len(data), 1)
|
|
self.assertEqual(data[0]["action_id"], "test-action-id")
|
|
|
|
def test_confirm_action_endpoint(self):
|
|
# Manually add a pending action
|
|
with patch('gui_2.uuid.uuid4', return_value="test-confirm-id"):
|
|
dialog = gui_2.ConfirmDialog("dir", ".")
|
|
self.app_instance._pending_actions[dialog._uid] = dialog
|
|
|
|
payload = {"approved": True}
|
|
response = self.client.post("/api/v1/confirm/test-confirm-id", json=payload, headers=self.headers)
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertTrue(dialog._done)
|
|
self.assertTrue(dialog._approved)
|
|
|
|
def test_list_sessions_endpoint(self):
|
|
# Ensure logs directory exists
|
|
Path("logs").mkdir(exist_ok=True)
|
|
# Create a dummy log
|
|
dummy_log = Path("logs/test_session_api.log")
|
|
dummy_log.write_text("dummy content")
|
|
|
|
try:
|
|
response = self.client.get("/api/v1/sessions", headers=self.headers)
|
|
self.assertEqual(response.status_code, 200)
|
|
data = response.json()
|
|
self.assertIn("test_session_api.log", data)
|
|
finally:
|
|
if dummy_log.exists():
|
|
dummy_log.unlink()
|
|
|
|
def test_get_context_endpoint(self):
|
|
response = self.client.get("/api/v1/context", headers=self.headers)
|
|
self.assertEqual(response.status_code, 200)
|
|
data = response.json()
|
|
self.assertIn("files", data)
|
|
self.assertIn("screenshots", data)
|
|
self.assertIn("files_base_dir", data)
|
|
|
|
def test_endpoint_no_api_key_configured(self):
|
|
# Test the security fix specifically
|
|
with patch.dict(self.app_instance.config, {"headless": {"api_key": ""}}):
|
|
response = self.client.get("/status", headers=self.headers)
|
|
self.assertEqual(response.status_code, 403)
|
|
self.assertEqual(response.json()["detail"], "API Key not configured on server")
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|