tests/test_start_server_script.py

from __future__ import annotations

import os
import shutil
import subprocess
import tempfile
import textwrap
from pathlib import Path
import unittest


class StartServerScriptTests(unittest.TestCase):
    def test_conda_env_uvicorn_is_preferred_on_macos_when_available(self):
        repo_root = Path(__file__).resolve().parent.parent
        source_script = repo_root / "start-server.sh"
        with tempfile.TemporaryDirectory() as tmp:
            temp_root = Path(tmp)
            shutil.copy(source_script, temp_root / "start-server.sh")
            os.chmod(temp_root / "start-server.sh", 0o755)

            fake_bin = temp_root / "bin"
            fake_bin.mkdir()
            fake_uname = fake_bin / "uname"
            fake_uname.write_text("#!/usr/bin/env bash\nprintf 'Darwin\\n'\n", encoding="utf-8")
            os.chmod(fake_uname, 0o755)

            conda_uvicorn = temp_root / ".conda-env" / "bin" / "uvicorn"
            conda_uvicorn.parent.mkdir(parents=True)
            conda_uvicorn.write_text(
                textwrap.dedent(
                    """\
                    #!/usr/bin/env bash
                    set -euo pipefail
                    printf 'UVICORN=conda\n'
                    printf 'ARGS=%s\n' "$*"
                    """
                ),
                encoding="utf-8",
            )
            os.chmod(conda_uvicorn, 0o755)

            venv_uvicorn = temp_root / ".venv" / "bin" / "uvicorn"
            venv_uvicorn.parent.mkdir(parents=True)
            venv_uvicorn.write_text(
                textwrap.dedent(
                    """\
                    #!/usr/bin/env bash
                    set -euo pipefail
                    printf 'UVICORN=venv\n'
                    """
                ),
                encoding="utf-8",
            )
            os.chmod(venv_uvicorn, 0o755)

            result = subprocess.run(
                [str(temp_root / "start-server.sh"), "dev"],
                cwd=temp_root,
                env={**os.environ, "PATH": f"{fake_bin}{os.pathsep}{os.environ['PATH']}"},
                capture_output=True,
                text=True,
                check=False,
            )

            self.assertEqual(result.returncode, 0, msg=result.stdout + result.stderr)
            self.assertIn("UVICORN=conda", result.stdout)
            self.assertNotIn("UVICORN=venv", result.stdout)
            self.assertIn("ARGS=app:app --host 127.0.0.1 --port 8000 --reload", result.stdout)

    def test_venv_uvicorn_is_preferred_on_linux_even_when_conda_env_exists(self):
        repo_root = Path(__file__).resolve().parent.parent
        source_script = repo_root / "start-server.sh"
        with tempfile.TemporaryDirectory() as tmp:
            temp_root = Path(tmp)
            shutil.copy(source_script, temp_root / "start-server.sh")
            os.chmod(temp_root / "start-server.sh", 0o755)

            fake_bin = temp_root / "bin"
            fake_bin.mkdir()
            fake_uname = fake_bin / "uname"
            fake_uname.write_text("#!/usr/bin/env bash\nprintf 'Linux\\n'\n", encoding="utf-8")
            os.chmod(fake_uname, 0o755)

            conda_uvicorn = temp_root / ".conda-env" / "bin" / "uvicorn"
            conda_uvicorn.parent.mkdir(parents=True)
            conda_uvicorn.write_text(
                textwrap.dedent(
                    """\
                    #!/usr/bin/env bash
                    set -euo pipefail
                    printf 'UVICORN=conda\n'
                    """
                ),
                encoding="utf-8",
            )
            os.chmod(conda_uvicorn, 0o755)

            venv_uvicorn = temp_root / ".venv" / "bin" / "uvicorn"
            venv_uvicorn.parent.mkdir(parents=True)
            venv_uvicorn.write_text(
                textwrap.dedent(
                    """\
                    #!/usr/bin/env bash
                    set -euo pipefail
                    printf 'UVICORN=venv\n'
                    printf 'ARGS=%s\n' "$*"
                    """
                ),
                encoding="utf-8",
            )
            os.chmod(venv_uvicorn, 0o755)

            result = subprocess.run(
                [str(temp_root / "start-server.sh"), "dev"],
                cwd=temp_root,
                env={**os.environ, "PATH": f"{fake_bin}{os.pathsep}{os.environ['PATH']}"},
                capture_output=True,
                text=True,
                check=False,
            )

            self.assertEqual(result.returncode, 0, msg=result.stdout + result.stderr)
            self.assertIn("UVICORN=venv", result.stdout)
            self.assertNotIn("UVICORN=conda", result.stdout)
            self.assertIn("ARGS=app:app --host 127.0.0.1 --port 8000 --reload", result.stdout)

    def test_model_weights_path_flag_is_forwarded_to_bucklakeai(self):
        repo_root = Path(__file__).resolve().parent.parent
        source_script = repo_root / "start-server.sh"
        with tempfile.TemporaryDirectory() as tmp:
            temp_root = Path(tmp)
            shutil.copy(source_script, temp_root / "start-server.sh")
            os.chmod(temp_root / "start-server.sh", 0o755)

            fake_uvicorn = temp_root / ".venv" / "bin" / "uvicorn"
            fake_uvicorn.parent.mkdir(parents=True)
            fake_uvicorn.write_text(
                textwrap.dedent(
                    """\
                    #!/usr/bin/env bash
                    set -euo pipefail
                    printf 'BL_MODEL_WEIGHTS_PATH=%s\n' "${BL_MODEL_WEIGHTS_PATH:-}"
                    printf 'BL_MODEL_VERSION=%s\n' "${BL_MODEL_VERSION:-}"
                    printf 'ARGS=%s\n' "$*"
                    """
                ),
                encoding="utf-8",
            )
            os.chmod(fake_uvicorn, 0o755)

            result = subprocess.run(
                [
                    str(temp_root / "start-server.sh"),
                    "--model-weights-path",
                    "/tmp/local/stock_prediction_model_anchored-path-v1_best.keras",
                    "dev",
                ],
                cwd=temp_root,
                capture_output=True,
                text=True,
                check=False,
            )

            self.assertEqual(result.returncode, 0, msg=result.stdout + result.stderr)
            self.assertIn(
                "BL_MODEL_WEIGHTS_PATH=/tmp/local/stock_prediction_model_anchored-path-v1_best.keras",
                result.stdout,
            )
            self.assertIn("BL_MODEL_VERSION=anchored-path-v1", result.stdout)
            self.assertIn("ARGS=app:app --host 127.0.0.1 --port 8000 --reload", result.stdout)

    def test_model_weights_path_flag_infers_anchored_path_v2(self):
        repo_root = Path(__file__).resolve().parent.parent
        source_script = repo_root / "start-server.sh"
        with tempfile.TemporaryDirectory() as tmp:
            temp_root = Path(tmp)
            shutil.copy(source_script, temp_root / "start-server.sh")
            os.chmod(temp_root / "start-server.sh", 0o755)

            fake_uvicorn = temp_root / ".venv" / "bin" / "uvicorn"
            fake_uvicorn.parent.mkdir(parents=True)
            fake_uvicorn.write_text(
                textwrap.dedent(
                    """\
                    #!/usr/bin/env bash
                    set -euo pipefail
                    printf 'BL_MODEL_WEIGHTS_PATH=%s\n' "${BL_MODEL_WEIGHTS_PATH:-}"
                    printf 'BL_MODEL_VERSION=%s\n' "${BL_MODEL_VERSION:-}"
                    """
                ),
                encoding="utf-8",
            )
            os.chmod(fake_uvicorn, 0o755)

            result = subprocess.run(
                [
                    str(temp_root / "start-server.sh"),
                    "--model-weights-path",
                    "/tmp/local/stock_prediction_model_anchored-path-v2_best.keras",
                    "dev",
                ],
                cwd=temp_root,
                capture_output=True,
                text=True,
                check=False,
            )

            self.assertEqual(result.returncode, 0, msg=result.stdout + result.stderr)
            self.assertIn("BL_MODEL_VERSION=anchored-path-v2", result.stdout)


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