CI + tests were added

Author: TejaMuraharirao

.github/workflows/ci.yml

@@ -0,0 +1,84 @@
+name: CI
+
+# Runs on every PR (and on pushes to the default branches, so a merge that
+# bypasses PR checks — e.g. an admin merge — still gets caught). See
+# ruff.toml for why the lint job's rule set is deliberately narrow, and
+# tests/conftest.py for why the test job needs a real Postgres service
+# rather than a mock.
+
+on:
+  pull_request:
+  push:
+    # TODO: remove "documents" once the CI setup itself has been verified —
+    # it's here temporarily so pushes to this feature branch (with no PR
+    # open yet) also trigger a run for testing.
+    branches: [master, main, documents]
+
+permissions:
+  contents: read
+
+jobs:
+  lint:
+    name: Ruff
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install dev dependencies
+        run: pip install -r requirements-dev.txt
+
+      - name: Ruff check
+        run: ruff check .
+
+  test:
+    name: Pytest
+    runs-on: ubuntu-latest
+
+    # tests/conftest.py requires a real PostgreSQL instance — it creates the
+    # aistudio_test database on first run and truncates tables between
+    # tests. A service container is the GitHub Actions equivalent of the
+    # `docker compose up -d postgres` step in local development.
+    services:
+      postgres:
+        image: postgres:15-alpine
+        env:
+          POSTGRES_USER: aistudio
+          POSTGRES_PASSWORD: aistudio
+          POSTGRES_DB: aistudio
+        ports:
+          - 5432:5432
+        options: >-
+          --health-cmd "pg_isready -U aistudio"
+          --health-interval 5s
+          --health-timeout 3s
+          --health-retries 5
+
+    env:
+      # Matches the "pytest directly on the host machine" case documented at
+      # the top of tests/conftest.py — the service container is reachable on
+      # localhost:5432 from the job's steps.
+      POSTGRES_HOST: localhost
+      POSTGRES_PORT: 5432
+      POSTGRES_USERNAME: aistudio
+      POSTGRES_PASSWORD: aistudio
+      POSTGRES_DATABASE: aistudio_test
+      RABBITMQ_URL: localhost
+      RABBITMQ_USERNAME: aistudio
+      RABBITMQ_PASSWORD: aistudio
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install dev dependencies
+        run: pip install -r requirements-dev.txt
+
+      - name: Run tests
+        run: pytest tests/ -v

app/main.py

@@ -3,6 +3,8 @@
 from fastapi.responses import JSONResponse
 from fastapi.exceptions import RequestValidationError
 
+from app.routers import system, ingest, benchmarks, results, jupyter, gpu_specs
+
 # The app instance
 app = FastAPI(
     title="AIStudio API",
@@ -60,8 +62,6 @@ async def generic_exception_handler(request: Request, exc: Exception):
         },
     )
 
-from app.routers import system, ingest, benchmarks, results, jupyter, gpu_specs
-
 # ── Routers ───────────────────────────────────────────────────────────────────
 app.include_router(system.router)
 app.include_router(ingest.router)

app/models/workload.py

@@ -3,7 +3,7 @@
 from datetime import datetime
 from typing import Optional
 
-from sqlalchemy import DateTime, Enum, ForeignKey, Integer, String, Text, func
+from sqlalchemy import DateTime, Enum, Integer, String, Text, func
 from sqlalchemy.dialects.postgresql import JSONB, UUID
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 

app/routers/benchmarks.py

@@ -79,6 +79,10 @@ async def get_benchmark_status(
 
     return {
         "task_id": workload.workload_id,
+        # Kept alongside task_id (same value) for callers that key off the
+        # DB column name directly — e.g. Composer lifecycle-event tracking,
+        # which correlates status polls against workload_id elsewhere.
+        "workload_id": workload.workload_id,
         "state": workload.state,
         "error_message": workload.error_message,
         "updated_at": workload.updated_at,

app/services/node_inspector.py

@@ -1,5 +1,4 @@
 from app.config import settings
-import json
 import logging
 from app.services.ssh_executor import SSHExecutor
 

app/services/ssh_executor.py

@@ -1,7 +1,7 @@
 import re
 import uuid
 import paramiko
-from typing import Generator, Optional
+from typing import Optional
 
 from app.database import SyncSessionLocal
 from app.models.task_log import TaskLog

app/utils/sse.py

@@ -8,14 +8,6 @@
 from datetime import datetime
 from typing import AsyncGenerator, Optional
 
-# Masks the last two octets of any IPv4 address in log lines so node IPs are
-# never exposed to the client. e.g. 10.6.12.26 → 10.6.x.x
-_IP_RE = re.compile(r'(?<![.\d])(\d{1,3}\.\d{1,3})\.\d{1,3}\.\d{1,3}(?![.\d])')
-
-
-def _mask_ip(text: str) -> str:
-    return _IP_RE.sub(r'\1.x.x', text)
-
 from fastapi import Request
 from sqlalchemy import select
 
@@ -24,6 +16,14 @@ def _mask_ip(text: str) -> str:
 from app.models.task_log import TaskLog
 from app.models.workload import Workload
 
+# Masks the last two octets of any IPv4 address in log lines so node IPs are
+# never exposed to the client. e.g. 10.6.12.26 → 10.6.x.x
+_IP_RE = re.compile(r'(?<![.\d])(\d{1,3}\.\d{1,3})\.\d{1,3}\.\d{1,3}(?![.\d])')
+
+
+def _mask_ip(text: str) -> str:
+    return _IP_RE.sub(r'\1.x.x', text)
+
 
 async def task_log_stream(
     workload_db_id: uuid.UUID,

pytest.ini

@@ -1,3 +1,9 @@
 [pytest]
 asyncio_mode = auto
+asyncio_default_fixture_loop_scope = session
+asyncio_default_test_loop_scope = session
 testpaths = tests
+; Safety net: without this, a hung async test (e.g. an SSE/streaming
+; endpoint that never terminates) blocks the whole suite indefinitely
+; instead of failing fast.
+timeout = 60

requirements-dev.txt

@@ -5,4 +5,6 @@
 
 pytest>=8.0.0
 pytest-asyncio>=0.24.0
+pytest-timeout>=2.3.0
 anyio[trio]>=4.0.0
+ruff>=0.16.0

ruff.toml

@@ -0,0 +1,41 @@
+# Ruff configuration — see .github/workflows/ci.yml for how this runs in CI.
+#
+# Scope is deliberately narrow: syntax errors and Pyflakes (undefined names,
+# unused imports, etc.) rather than Ruff's full default rule set (which also
+# pulls in pyupgrade, flake8-bugbear, flake8-datetimez and a dozen other
+# plugins). Turning all of that on for a codebase that has never been linted
+# produces hundreds of pre-existing findings unrelated to any given PR, which
+# makes the check noise rather than a real gate. Start narrow and correct;
+# widen deliberately later if the team wants stricter style enforcement.
+target-version = "py311"
+
+[lint]
+# E4/E7/E9 = pycodestyle error subset (real mistakes, not style nits).
+# F       = Pyflakes (undefined names, unused imports, etc).
+select = ["E4", "E7", "E9", "F"]
+
+[lint.per-file-ignores]
+# Models are imported here purely so Alembic's autogenerate can discover them
+# via `Base.metadata` — the names are never referenced directly in this file.
+"app/models/__init__.py" = ["F401"]
+
+# SQLAlchemy 2.0 `Mapped["ClassName"]` relationship annotations use forward
+# references as plain strings resolved by SQLAlchemy at mapper-configuration
+# time, not by Python's own name resolution — Pyflakes has no way to know
+# that and flags them as undefined names. Known false positive; see
+# https://github.com/astral-sh/ruff/issues/8811.
+"app/models/*.py" = ["F821"]
+
+# alembic/env.py intentionally star-imports every model so Alembic's
+# autogenerate can see them via Base.metadata; Pyflakes can't verify names
+# from a star import.
+"alembic/env.py" = ["F403"]
+
+# Historical, already-applied migrations are immutable — don't hand-edit a
+# shipped revision just to satisfy a linter added long after it merged.
+"alembic/versions/*.py" = ["F401"]
+
+# conftest.py must set POSTGRES_*/RABBITMQ_* environment variables BEFORE
+# importing app.config, since pydantic-settings reads env vars at class-body
+# time. That requires imports after statements — intentional, not an oversight.
+"tests/conftest.py" = ["E402"]