JQC_multi_tenant/scripts/scratch_bootstrap_test.py

#!/usr/bin/env python3
"""
scratch_bootstrap_test.py
=========================
One-shot, self-cleaning test that proves the squashed baseline DDL executes on
real MySQL and that bootstrap_tenant() builds a complete, head-stamped schema.

It mirrors what MT-3 provisioning will do:
  1. (as a MySQL admin) CREATE a throwaway database + user + grant
  2. insert a temporary tenant row in the control DB (encrypted creds)
  3. run control.tenant_migrate.bootstrap_tenant() against it
  4. verify: all expected tables built + alembic_version stamped at chain head
  5. tear everything down (tenant row, provisioning jobs, database, user)

Nothing touches your real LT database or any real tenant. Run from the REPO ROOT.

Prerequisites in the environment (same as the runner):
    export CONTROL_DATABASE_URL='mysql+pymysql://jqc_control:pw@127.0.0.1/jqc_control'
    export CONTROL_FERNET_KEY='...'

Example:
    python scratch_bootstrap_test.py \
        --admin-user root --admin-password 'ROOTPW' \
        --host 127.0.0.1 --port 3306

Optional overrides: --test-db, --test-user, --test-password, --user-host, --keep
"""

import argparse
import sys

import pymysql
from sqlalchemy import create_engine, inspect, text

from control.base import control_session
from control.models import Plan, Tenant, ProvisioningJob
from control.seed import seed_plans
from control.tenant_migrate import bootstrap_tenant, chain_head

EXPECTED_TABLES = {
    'api_device_tokens', 'api_refresh_tokens', 'areas', 'audit_logs', 'broadcasts',
    'checklist_items', 'customer_assignments', 'device_registrations', 'facilities',
    'facility_score_alerts', 'inspection_results', 'inspection_templates',
    'inspections', 'inspector_assignments', 'issue_comments', 'issue_followers',
    'issues', 'notification_matrix', 'notification_preferences', 'notifications',
    'projects', 'scheduled_reports', 'support_ticket_replies', 'support_tickets',
    'users',
}


def _admin_conn(args):
    return pymysql.connect(
        host=args.host, port=args.port,
        user=args.admin_user, password=args.admin_password,
        autocommit=True,
    )


def _admin_exec(conn, sql):
    with conn.cursor() as cur:
        cur.execute(sql)


def create_test_db(args):
    conn = _admin_conn(args)
    try:
        _admin_exec(conn, f"DROP DATABASE IF EXISTS `{args.test_db}`")
        _admin_exec(conn, f"CREATE DATABASE `{args.test_db}` "
                          f"CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
        _admin_exec(conn, f"DROP USER IF EXISTS '{args.test_user}'@'{args.user_host}'")
        _admin_exec(conn, f"CREATE USER '{args.test_user}'@'{args.user_host}' "
                          f"IDENTIFIED BY '{args.test_password}'")
        _admin_exec(conn, f"GRANT ALL PRIVILEGES ON `{args.test_db}`.* "
                          f"TO '{args.test_user}'@'{args.user_host}'")
        _admin_exec(conn, "FLUSH PRIVILEGES")
    finally:
        conn.close()
    print(f"  [setup] created database `{args.test_db}` + user "
          f"'{args.test_user}'@'{args.user_host}'")


def drop_test_db(args):
    conn = _admin_conn(args)
    try:
        _admin_exec(conn, f"DROP DATABASE IF EXISTS `{args.test_db}`")
        _admin_exec(conn, f"DROP USER IF EXISTS '{args.test_user}'@'{args.user_host}'")
        _admin_exec(conn, "FLUSH PRIVILEGES")
    finally:
        conn.close()
    print(f"  [cleanup] dropped database + user")


def insert_temp_tenant(args):
    seed_plans()  # idempotent — ensures at least one plan exists
    with control_session() as s:
        plan = s.query(Plan).order_by(Plan.id).first()
        t = Tenant(
            slug=args.test_slug, name='SCRATCH TEST (temporary)',
            plan_id=plan.id, status='provisioning',
            db_host=args.host, db_port=args.port,
            db_name=args.test_db, db_user=args.test_user,
        )
        t.set_db_password(args.test_password)
        s.add(t)
        s.flush()
        tid = t.id
    print(f"  [setup] inserted temp tenant id={tid} slug='{args.test_slug}'")
    return tid


def delete_temp_tenant(tid):
    with control_session() as s:
        s.query(ProvisioningJob).filter_by(tenant_id=tid).delete()
        t = s.get(Tenant, tid)
        if t is not None:
            s.delete(t)
    print(f"  [cleanup] removed temp tenant id={tid} + its provisioning jobs")


def verify(args):
    uri = (f"mysql+pymysql://{args.test_user}:{args.test_password}"
           f"@{args.host}:{args.port}/{args.test_db}")
    engine = create_engine(uri, future=True)
    ok = True
    try:
        insp = inspect(engine)
        built = set(insp.get_table_names())
        built.discard('alembic_version')

        missing = EXPECTED_TABLES - built
        extra = built - EXPECTED_TABLES
        print(f"  tables built: {len(built)} (expected {len(EXPECTED_TABLES)})")
        if missing:
            print(f"  !! MISSING tables: {sorted(missing)}"); ok = False
        if extra:
            print(f"  ?? unexpected extra tables: {sorted(extra)}")

        with engine.connect() as c:
            ver = c.execute(text("SELECT version_num FROM alembic_version")).scalar()
        head = chain_head()
        print(f"  alembic_version: {ver}  (chain head: {head})")
        if ver != head:
            print("  !! alembic_version is not at chain head"); ok = False

        # spot-check a couple of tables have their columns
        ucols = {c['name'] for c in insp.get_columns('users')}
        for need in ('id', 'username', 'email', 'password_hash', 'role'):
            if need not in ucols:
                print(f"  !! users.{need} missing"); ok = False
        print(f"  users columns: {len(ucols)} (incl id/username/email/role)")
    finally:
        engine.dispose()
    return ok


def main():
    ap = argparse.ArgumentParser(description='Scratch MySQL bootstrap test')
    ap.add_argument('--admin-user', required=True, help='MySQL admin user (can CREATE DB/USER)')
    ap.add_argument('--admin-password', required=True)
    ap.add_argument('--host', default='127.0.0.1')
    ap.add_argument('--port', type=int, default=3306)
    ap.add_argument('--test-db', default='jqc_scratch_test')
    ap.add_argument('--test-user', default='jqc_scratch_user')
    ap.add_argument('--test-password', default='Scratch_pw_123!')
    ap.add_argument('--test-slug', default='scratchtest')
    ap.add_argument('--user-host', default='%', help="MySQL host part for the test user")
    ap.add_argument('--keep', action='store_true', help='Skip cleanup (inspect manually)')
    args = ap.parse_args()

    print("== JQC scratch bootstrap test ==")
    tid = None
    passed = False
    try:
        create_test_db(args)
        tid = insert_temp_tenant(args)
        print("  [run] bootstrap_tenant() ...")
        with control_session() as s:
            t = s.get(Tenant, tid)
            ref = type('R', (), {'id': t.id, 'db_uri': t.db_uri})()
        applied = bootstrap_tenant(ref)
        print(f"  [run] bootstrapped -> {applied}")
        passed = verify(args)
    except Exception as e:
        print(f"  !! ERROR: {type(e).__name__}: {e}")
        passed = False
    finally:
        if not args.keep:
            if tid is not None:
                try:
                    delete_temp_tenant(tid)
                except Exception as e:
                    print(f"  !! tenant cleanup failed: {e}")
            try:
                drop_test_db(args)
            except Exception as e:
                print(f"  !! db cleanup failed: {e}")
        else:
            print("  [keep] left test DB + tenant in place for inspection")

    print()
    print("RESULT:", "PASS ✅" if passed else "FAIL ❌")
    sys.exit(0 if passed else 1)


if __name__ == '__main__':
    main()