import uuid from fastapi.testclient import TestClient from sqlalchemy import delete from app.core.security import create_access_token, hash_password from app.db.session import SessionLocal from app.main import app from app.models.branch import Branch from app.models.business import Business from app.models.membership import Membership, MembershipBranch from app.models.user import User from app.services.rbac import require_branch_access def headers(user): return {"Authorization": f"Bearer {create_access_token(str(user.id))}"} def setup(): db = SessionLocal(); marker = uuid.uuid4().hex ceo = User(full_name="CEO", email=f"ceo-{marker}@example.com", password_hash=hash_password("TestPassword123!"), is_email_verified=True) admin = User(full_name="Admin", email=f"admin-{marker}@example.com", password_hash=hash_password("TestPassword123!"), is_email_verified=True) cashier = User(full_name="Cashier", email=f"cashier-{marker}@example.com", password_hash=hash_password("TestPassword123!"), is_email_verified=True) business = Business(name="Staff Store") db.add_all([ceo, admin, cashier, business]); db.flush() main = Branch(business_id=business.id, name="Main", is_main=True) other = Branch(business_id=business.id, name="Other") ceo_member = Membership(business_id=business.id, user_id=ceo.id, role="CEO", all_branches=True) admin_member = Membership(business_id=business.id, user_id=admin.id, role="ADMIN", all_branches=True) cashier_member = Membership(business_id=business.id, user_id=cashier.id, role="CASHIER") db.add_all([main, other, ceo_member, admin_member, cashier_member]); db.commit() return db, business, main, other, ceo, admin, cashier, ceo_member, admin_member, cashier_member def cleanup(db, business, users): db.execute(delete(MembershipBranch).where(MembershipBranch.membership_id.in_(select_memberships(business.id)))) db.execute(delete(Membership).where(Membership.business_id == business.id)) db.execute(delete(Branch).where(Branch.business_id == business.id)) db.execute(delete(Business).where(Business.id == business.id)) db.execute(delete(User).where(User.id.in_([user.id for user in users]))) db.commit(); db.close() def select_memberships(business_id): from sqlalchemy import select return select(Membership.id).where(Membership.business_id == business_id) def test_staff_rbac_and_branch_access(): db, business, main, other, ceo, admin, cashier, ceo_m, admin_m, cashier_m = setup() try: with TestClient(app) as client: listed = client.get(f"/api/businesses/{business.id}/staff", headers=headers(ceo)) assert listed.status_code == 200 and len(listed.json()) == 3 assert client.get(f"/api/businesses/{business.id}/staff").status_code == 401 assert client.get(f"/api/businesses/{business.id}/staff", headers=headers(cashier)).status_code == 403 changed = client.patch(f"/api/businesses/{business.id}/staff/{cashier_m.id}/role", json={"role": "ACCOUNTANT"}, headers=headers(ceo)) assert changed.status_code == 200 and changed.json()["role"] == "ACCOUNTANT" assert client.patch(f"/api/businesses/{business.id}/staff/{cashier_m.id}/role", json={"role": "CEO"}, headers=headers(admin)).status_code == 422 assert client.patch(f"/api/businesses/{business.id}/staff/{admin_m.id}/role", json={"role": "CEO"}, headers=headers(admin)).status_code == 422 denied = client.patch(f"/api/businesses/{business.id}/staff/{ceo_m.id}/status", json={"is_active": False}, headers=headers(ceo)) assert denied.status_code == 422 denied = client.patch(f"/api/businesses/{business.id}/staff/{ceo_m.id}/role", json={"role": "ADMIN"}, headers=headers(ceo)) assert denied.status_code == 422 assigned = client.put(f"/api/businesses/{business.id}/staff/{cashier_m.id}/branches", json={"all_branches": False, "branch_ids": [str(main.id)]}, headers=headers(ceo)) assert assigned.status_code == 200 db.refresh(cashier_m) assert require_branch_access(db, cashier_m, business.id, main.id).id == main.id try: require_branch_access(db, cashier_m, business.id, other.id) assert False, "unassigned branch was allowed" except Exception as exc: assert getattr(exc, "status_code", None) == 403 duplicate = client.put(f"/api/businesses/{business.id}/staff/{cashier_m.id}/branches", json={"all_branches": False, "branch_ids": [str(main.id), str(main.id)]}, headers=headers(ceo)) assert duplicate.status_code == 422 all_access = client.put(f"/api/businesses/{business.id}/staff/{cashier_m.id}/branches", json={"all_branches": True, "branch_ids": []}, headers=headers(ceo)) assert all_access.status_code == 200 db.refresh(cashier_m) assert require_branch_access(db, cashier_m, business.id, other.id).id == other.id assert client.patch(f"/api/businesses/{business.id}/staff/{cashier_m.id}/status", json={"is_active": False}, headers=headers(ceo)).status_code == 200 assert client.get(f"/api/businesses/{business.id}/staff", headers=headers(cashier)).status_code == 404 finally: cleanup(db, business, [ceo, admin, cashier]) def test_staff_business_isolation_and_registered_user_addition(): db, business, main, other, ceo, admin, cashier, ceo_m, admin_m, cashier_m = setup() second = Business(name="Other Business") db.add(second); db.flush() second_branch = Branch(business_id=second.id, name="Other Main") second_member = Membership(business_id=second.id, user_id=cashier.id, role="CASHIER") db.add_all([second_branch, second_member]); db.commit() try: with TestClient(app) as client: # A membership ID from another business cannot be read or modified. assert client.get(f"/api/businesses/{business.id}/staff/{second_member.id}", headers=headers(ceo)).status_code == 404 assert client.put(f"/api/businesses/{business.id}/staff/{second_member.id}/branches", json={"all_branches": True}, headers=headers(ceo)).status_code == 404 # Existing, registered users can be safely attached once; duplicate attachment is rejected. newcomer = User(full_name="New", email=f"new-{uuid.uuid4().hex}@example.com", password_hash=hash_password("TestPassword123!"), is_email_verified=True) db.add(newcomer); db.commit() payload = {"email": newcomer.email, "role": "STOREKEEPER", "all_branches": False, "branch_ids": [str(main.id)]} added = client.post(f"/api/businesses/{business.id}/staff", json=payload, headers=headers(ceo)) assert added.status_code == 201 assert client.post(f"/api/businesses/{business.id}/staff", json=payload, headers=headers(ceo)).status_code == 422 db.execute(delete(User).where(User.id == newcomer.id)); db.commit() finally: db.execute(delete(MembershipBranch).where(MembershipBranch.membership_id.in_(select_memberships(second.id)))) db.execute(delete(Membership).where(Membership.business_id == second.id)) db.execute(delete(Branch).where(Branch.business_id == second.id)) db.execute(delete(Business).where(Business.id == second.id)); db.commit() cleanup(db, business, [ceo, admin, cashier])