sales-trainer/backend/scripts/test_saas_tenant.py

"""Test: SaaS multi-tenant isolation + hardening basics."""
import os, sys, tempfile, warnings
from pathlib import Path

warnings.filterwarnings("ignore")
BACKEND = str(Path(__file__).resolve().parents[1])
sys.path.insert(0, BACKEND)

from app.factory import create_app
from app.config import Config

td = tempfile.mkdtemp()
Config.DATA_DIR = Path(td)
sys.path.insert(0, BACKEND + "/scripts")
from mock_llm import MockLLM

app = create_app()
app.extensions["llm"] = MockLLM()
C = app.test_client()

def tok(u, p): return C.post("/api/auth/login", json={"username": u, "password": p}).get_json()["token"]

# super admin setup
AT = tok("admin", "1234"); AH = {"Authorization": f"Bearer {AT}"}
C.post("/api/auth/setup", headers=AH, json={"username": "admin", "email": "a@b.co", "password": "newpass", "accepted_terms": True})
AT = tok("admin", "newpass"); AH = {"Authorization": f"Bearer {AT}"}

# create org1 data (a group)
gid = C.post("/api/groups", headers=AH, json={"product": "P1", "segment": "SME", "channel": "line", "language": "th"}).get_json()["group"]["id"]
assert gid
print("[ok] super_admin created group in default org")

# create a brand-new org (org2) with its own admin via new_org
r = C.post("/api/admin/users", headers=AH, json={
    "username": "adm2", "name": "Org2 Admin", "password": "pppp", "role": "admin", "new_org": True
})
assert r.status_code == 201, r.get_json()
org2_id = r.get_json()["org_id"]
assert org2_id and org2_id != "org-default"
print("[ok] super_admin created new org (id=%s) with its admin" % org2_id[:8])

# org2 admin login + try to read org1's group -> must fail
AT2 = tok("adm2", "pppp"); AH2 = {"Authorization": f"Bearer {AT2}"}
r = C.get(f"/api/groups/{gid}", headers=AH2)
assert r.status_code == 403, ("org2 admin should NOT read org1 group", r.get_json())
print("[ok] org2 admin blocked from org1 group (403)")

# org2 admin list groups -> sees none of org1's (empty)
r = C.get("/api/groups", headers=AH2)
assert r.status_code == 200
groups = r.get_json()["groups"]
assert all(g.get("id") != gid for g in groups), "org2 sees org1's group"
print("[ok] org2 admin cannot see org1 group in list")

# org2 admin list users -> only org2 users (adm2), not the super admin
r = C.get("/api/admin/users", headers=AH2)
names = [u.get("username") for u in r.get_json()["users"]]
assert "adm2" in names and "admin" not in names, names
print("[ok] org2 admin lists only org2 users")

# super_admin platform orgs view sees both orgs
r = C.get("/api/admin/orgs", headers=AH)
orgs = r.get_json()["orgs"]
assert len(orgs) >= 2, orgs
print("[ok] super_admin sees all orgs (%d)" % len(orgs))

# org2 admin sees only own org in /orgs
r = C.get("/api/admin/orgs", headers=AH2)
orgs = r.get_json()["orgs"]
assert len(orgs) == 1 and orgs[0]["id"] == org2_id, orgs
print("[ok] org2 admin sees only its own org")

# Rate limiting: reset mem then hammer login on a fake user -> 429 lockout on user key
from app.services import rate_limit
rate_limit._mem.clear()
code = None
for _ in range(11):
    code = C.post("/api/auth/login", json={"username": "nobody", "password": "x"}).status_code
assert code == 429, ("expected 429 after login hammering", code)
rate_limit._mem.clear()
print("[ok] login rate-limit returns 429 after abuse")

# ── Phase 3: seat enforcement + inactive org gating + plan update ──────────────
# default org (org-default) was created with seats=5, plan=trial, active=True
org = C.get("/api/admin/orgs", headers=AH).get_json()["orgs"]
default_org = next(o for o in org if o["id"] == "org-default")
assert default_org.get("plan") == "trial" and default_org.get("seats") == 5 and default_org.get("active") is True
print("[ok] default org has plan=trial, seats=5, active=True")

# seat limit: fill the default org up to seats=5 -> 6th create must fail
# dynamic: count current users in default org, then try to add enough to exceed seats
current_count = len([u for u in C.get("/api/admin/users", headers=AH).get_json()["users"] if u.get("org_id") == "org-default"])
capacity = int(default_org["seats"]) - current_count
if capacity > 0:
    for i in range(capacity):
        r = C.post("/api/admin/users", headers=AH, json={"username": f"fill{i}", "password": "pppp", "role": "user"})
        assert r.status_code == 201, r.get_json()
# now one more must exceed seats
r = C.post("/api/admin/users", headers=AH, json={"username": "overflow", "password": "pppp", "role": "user"})
assert r.status_code == 400, ("seat limit should block", r.get_json())
print("[ok] seat limit blocks extra user")

# super_admin can bump seats then create succeeds
org_id_def = "org-default"
C.patch(f"/api/admin/orgs/{org_id_def}", headers=AH, json={"seats": 99})
r = C.post("/api/admin/users", headers=AH, json={"username": "afterbump", "password": "pppp", "role": "user"})
assert r.status_code == 201, r.get_json()
print("[ok] super_admin can raise seats then add user")

# deactivate org -> login blocked for its users
C.patch(f"/api/admin/orgs/{org_id_def}", headers=AH, json={"active": False})
code = C.post("/api/auth/login", json={"username": "admin", "password": "newpass"}).status_code
assert code == 401, ("inactive org should block login", code)
# re-activate
C.patch(f"/api/admin/orgs/{org_id_def}", headers=AH, json={"active": True})
print("[ok] inactive org blocks login; re-activate restores")

# plain admin cannot patch orgs (403)
r = C.patch(f"/api/admin/orgs/{org_id_def}", headers=AH2, json={"seats": 200})
assert r.status_code == 403, r.status_code
print("[ok] plain admin cannot update org plan")

print("ALL SAAS MULTI-TENANT TESTS PASSED")