Google Antigravity

"""Launch contract regressions: slow diagnostics must not break the dock."""
import json
import os
from pathlib import Path
import signal
import socket
import subprocess
import sys
import tempfile
import threading
import time
import unittest
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from unittest.mock import patch

import adom_antigravity as app


class LegacyHandler(BaseHTTPRequestHandler):
    def log_message(self, *args):
        pass

    def do_GET(self):
        if self.path != '/api/status':
            self.send_error(404)
            return
        time.sleep(1.6)  # ConfRoomROG; longer than both original startup budgets.
        body = json.dumps({'app': 'adom/antigravity', 'version': app.VERSION}).encode()
        self.send_response(200)
        self.send_header('Content-Length', str(len(body)))
        self.end_headers()
        self.wfile.write(body)


class StartupTests(unittest.TestCase):
    @staticmethod
    def free_port():
        with socket.socket() as sock:
            sock.bind(('127.0.0.1', 0))
            return sock.getsockname()[1]

    def serve_fixture(self, handler):
        server = ThreadingHTTPServer(('127.0.0.1', 0), handler)
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        self.addCleanup(server.server_close)
        self.addCleanup(server.shutdown)
        return server.server_port

    def launch(self, port, home):
        return subprocess.run(
            [sys.executable, str(Path(app.__file__)), 'serve', '--print-url', '--port', str(port)],
            env={**os.environ, 'HOME': home}, capture_output=True, text=True, timeout=25)

    def test_legacy_slow_status_is_reused_without_spawning(self):
        port = self.serve_fixture(LegacyHandler)
        with tempfile.TemporaryDirectory() as home:
            result = self.launch(port, home)
            self.assertEqual(result.returncode, 0, result.stderr)
            self.assertEqual(result.stdout.strip(), f'http://127.0.0.1:{port}/')
            self.assertFalse((Path(home) / '.local/state/adom/antigravity/dashboard.log').exists())

    def test_health_does_not_run_external_diagnostics(self):
        port = self.serve_fixture(app.Handler)
        with patch.object(app, 'context', side_effect=AssertionError('diagnostics ran')):
            self.assertEqual(app.server_identity(f'http://127.0.0.1:{port}/')['pid'], os.getpid())

    def test_cold_start_then_relaunch_keeps_same_process(self):
        with socket.socket() as sock:
            sock.bind(('127.0.0.1', 0))
            port = sock.getsockname()[1]
        with tempfile.TemporaryDirectory() as home:
            first = self.launch(port, home)
            self.assertEqual(first.returncode, 0, first.stderr)
            url = first.stdout.strip()
            pid = app.server_identity(url)['pid']
            try:
                second = self.launch(port, home)
                self.assertEqual(second.returncode, 0, second.stderr)
                self.assertEqual(second.stdout.strip(), url)
                self.assertEqual(app.server_identity(url)['pid'], pid)
            finally:
                os.kill(pid, signal.SIGTERM)

    def test_other_service_is_not_terminated(self):
        class OtherHandler(LegacyHandler):
            def do_GET(self):
                body = b'{"app":"another-app","version":"1"}'
                self.send_response(200)
                self.send_header('Content-Length', str(len(body)))
                self.end_headers()
                self.wfile.write(body)
        port = self.serve_fixture(OtherHandler)
        with tempfile.TemporaryDirectory() as home:
            result = self.launch(port, home)
            self.assertNotEqual(result.returncode, 0)
            self.assertEqual(app.server_identity(f'http://127.0.0.1:{port}/')['app'], 'another-app')

    @unittest.skipUnless(os.environ.get('ANTIGRAVITY_PREDECESSOR'), 'requires a registry predecessor')
    def test_registry_predecessor_upgrade_is_scoped_to_port(self):
        self.check_upgrade(os.environ['ANTIGRAVITY_PREDECESSOR'])

    def test_upgrade_of_server_with_current_lifecycle(self):
        with tempfile.TemporaryDirectory() as directory:
            predecessor = Path(directory) / Path(app.__file__).name
            predecessor.write_text(Path(app.__file__).read_text().replace(
                "VERSION = '0.1.6'", "VERSION = '0.1.5'"))
            self.check_upgrade(str(predecessor))

    def check_upgrade(self, predecessor):
        with tempfile.TemporaryDirectory() as home:
            ports = [self.free_port(), self.free_port()]
            old = [subprocess.Popen(
                [sys.executable, predecessor, '_serve', '--port', str(port)],
                env={**os.environ, 'HOME': home}, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                for port in ports]
            current_pid = None
            try:
                for index, port in enumerate(ports):
                    for _ in range(50):
                        self.assertIsNone(old[index].poll(), 'Predecessor process exited')
                        try:
                            state = app.server_identity(f'http://127.0.0.1:{port}/')
                            self.assertEqual(state['version'], '0.1.5')
                            break
                        except OSError:
                            time.sleep(.1)
                    else:
                        self.fail('Predecessor never started')
                result = self.launch(ports[0], home)
                self.assertEqual(result.returncode, 0, result.stderr)
                state = app.server_identity(result.stdout.strip())
                current_pid = state['pid']
                self.assertEqual(state['version'], app.VERSION)
                old[0].wait(timeout=5)
                self.assertIsNone(old[1].poll(), 'Another port was affected')
                again = self.launch(ports[0], home)
                self.assertEqual(again.returncode, 0, again.stderr)
                self.assertEqual(app.server_identity(again.stdout.strip())['pid'], current_pid)
            finally:
                if current_pid:
                    os.kill(current_pid, signal.SIGTERM)
                for process in old:
                    if process.poll() is None:
                        process.terminate()
                    process.wait(timeout=5)


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