163 lines
7.8 KiB
Python
163 lines
7.8 KiB
Python
import asyncio
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
|
|
import httpx
|
|
from fastapi import FastAPI, File, Request, UploadFile
|
|
|
|
from backend.app import db, main
|
|
from backend.app.config import settings
|
|
from backend.app.request_limits import InstallationBodyLimitMiddleware
|
|
from backend.app.services import setup
|
|
|
|
|
|
class InstallationLifecycleTests(unittest.IsolatedAsyncioTestCase):
|
|
def setUp(self):
|
|
temporary = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
|
|
self.addCleanup(temporary.cleanup)
|
|
patches = [
|
|
patch.object(settings, "sqlite_path", str(Path(temporary.name) / "magent.db")),
|
|
patch.object(settings, "jwt_secret", "installation-lifecycle-secret-1234567890"),
|
|
patch.object(settings, "settings_encryption_key", None),
|
|
patch.object(settings, "admin_password", ""),
|
|
patch.object(settings, "setup_token", "operator-setup-token-at-least-32-characters"),
|
|
patch.object(main, "_background_started", False),
|
|
patch.object(main, "_background_tasks", []),
|
|
patch.object(main, "start_metrics"),
|
|
patch.object(main, "configure_logging"),
|
|
patch.dict("os.environ", {"BACKGROUND_TASKS_ENABLED": "true"}),
|
|
]
|
|
for item in patches:
|
|
item.start()
|
|
self.addCleanup(item.stop)
|
|
|
|
async def test_fresh_start_waits_for_admin_and_completion_then_starts_workers_once(self):
|
|
with patch.object(main, "_launch_background_task") as launch:
|
|
await main.startup()
|
|
self.assertEqual(setup.get_public_setup_status(), {"setup_required": True, "needs_admin": True})
|
|
launch.assert_not_called()
|
|
setup.bootstrap_administrator(settings.setup_token, "owner", "new-password-12345")
|
|
await main._start_background_tasks()
|
|
launch.assert_not_called()
|
|
setup.complete_setup()
|
|
await main.app.state.on_setup_complete()
|
|
await main.app.state.on_setup_complete()
|
|
self.assertEqual(launch.call_count, 9)
|
|
|
|
async def test_upgraded_install_starts_normally_without_setup_token(self):
|
|
db.init_db()
|
|
db.create_user("owner", "existing-password-12345", role="admin")
|
|
settings.setup_token = ""
|
|
with patch.object(main, "_launch_background_task") as launch:
|
|
await main.startup()
|
|
self.assertFalse(setup.is_setup_required())
|
|
self.assertEqual(launch.call_count, 9)
|
|
|
|
async def test_disabled_workers_stay_disabled_after_setup(self):
|
|
setup.initialize_setup_state()
|
|
db.init_db()
|
|
setup.bootstrap_administrator(settings.setup_token, "owner", "new-password-12345")
|
|
setup.complete_setup()
|
|
with patch.dict("os.environ", {"BACKGROUND_TASKS_ENABLED": "false"}), patch.object(main, "_launch_background_task") as launch:
|
|
await main._start_background_tasks()
|
|
launch.assert_not_called()
|
|
|
|
async def test_bad_secret_stops_before_restore_or_database_initialization(self):
|
|
settings.jwt_secret = "short"
|
|
with patch.object(main, "apply_pending_restore") as restore, patch.object(main, "init_db") as initialize:
|
|
with self.assertRaisesRegex(RuntimeError, "JWT_SECRET"):
|
|
await main.startup()
|
|
restore.assert_not_called()
|
|
initialize.assert_not_called()
|
|
|
|
async def test_restore_failure_stops_before_initialization_and_workers(self):
|
|
with patch.object(main, "apply_pending_restore", side_effect=RuntimeError("restore failed")), patch.object(main, "init_db") as initialize, patch.object(main, "_launch_background_task") as launch:
|
|
with self.assertRaisesRegex(RuntimeError, "restore failed"):
|
|
await main.startup()
|
|
initialize.assert_not_called()
|
|
launch.assert_not_called()
|
|
|
|
async def test_startup_order_is_restore_then_setup_marker_then_schema(self):
|
|
calls = Mock()
|
|
calls.attach_mock(Mock(wraps=main.apply_pending_restore), "restore")
|
|
calls.attach_mock(Mock(wraps=main.initialize_setup_state), "setup")
|
|
calls.attach_mock(Mock(wraps=main.init_db), "schema")
|
|
with patch.object(main, "apply_pending_restore", calls.restore), patch.object(main, "initialize_setup_state", calls.setup), patch.object(main, "init_db", calls.schema):
|
|
await main.startup()
|
|
self.assertEqual([call[0] for call in calls.mock_calls], ["restore", "setup", "schema"])
|
|
|
|
def test_missing_token_does_not_allow_fresh_bootstrap(self):
|
|
setup.initialize_setup_state()
|
|
db.init_db()
|
|
settings.setup_token = ""
|
|
with self.assertRaisesRegex(RuntimeError, "SETUP_TOKEN"):
|
|
main._enforce_secure_startup_configuration()
|
|
|
|
def test_destination_environment_does_not_add_an_admin_to_restored_accounts(self):
|
|
db.init_db()
|
|
db.create_user("restored-owner", "existing-password-12345", role="admin")
|
|
with patch.object(settings, "admin_username", "host-bootstrap"), patch.object(settings, "admin_password", "new-host-password-12345"):
|
|
db.init_db()
|
|
self.assertIsNone(db.get_user_by_username("host-bootstrap"))
|
|
|
|
async def test_shutdown_cancels_workers_and_allows_next_start(self):
|
|
task = asyncio.create_task(asyncio.Event().wait())
|
|
main._background_tasks.append(task)
|
|
main._background_started = True
|
|
await main.shutdown()
|
|
self.assertTrue(task.cancelled())
|
|
self.assertEqual(main._background_tasks, [])
|
|
self.assertFalse(main._background_started)
|
|
|
|
|
|
class InstallationRequestLimitsTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_rejects_oversized_declared_body_before_parser(self):
|
|
app = FastAPI()
|
|
app.add_middleware(InstallationBodyLimitMiddleware)
|
|
|
|
@app.post("/setup/bootstrap")
|
|
async def bootstrap(request: Request):
|
|
self.fail("Body must be rejected before the endpoint")
|
|
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
response = await client.post("/setup/bootstrap", content=b"{}", headers={"Content-Length": "999999"})
|
|
self.assertEqual(response.status_code, 413)
|
|
|
|
async def test_counts_chunks_with_missing_or_forged_content_length(self):
|
|
app = FastAPI()
|
|
app.add_middleware(InstallationBodyLimitMiddleware)
|
|
|
|
@app.post("/setup/bootstrap")
|
|
async def bootstrap(request: Request):
|
|
return await request.json()
|
|
|
|
async def chunks():
|
|
yield b'{"token":"'
|
|
yield b"a" * 17000
|
|
yield b'"}'
|
|
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
for headers in ({}, {"Content-Length": "1"}):
|
|
response = await client.post("/setup/bootstrap", content=chunks(), headers=headers)
|
|
self.assertEqual(response.status_code, 413)
|
|
|
|
async def test_multipart_stream_limit_is_413_not_parser_500(self):
|
|
app = FastAPI()
|
|
app.add_middleware(InstallationBodyLimitMiddleware)
|
|
|
|
@app.post("/admin/backups/restore")
|
|
async def restore(file: UploadFile = File(...)):
|
|
return {"size": file.size}
|
|
|
|
async def chunks():
|
|
yield b'--boundary\r\nContent-Disposition: form-data; name="file"; filename="backup"\r\n\r\n'
|
|
yield b"a" * 2048
|
|
yield b"\r\n--boundary--\r\n"
|
|
|
|
with patch("backend.app.request_limits.RESTORE_BODY_LIMIT", 1024):
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
response = await client.post("/admin/backups/restore", content=chunks(), headers={"Content-Type": "multipart/form-data; boundary=boundary"})
|
|
self.assertEqual(response.status_code, 413)
|