"""Linking a GP entity to its investor account (so Assets can pull real balances)."""

from tests.conftest import make_user
from ten31portal.models import Entity, EntityType, UserRole


def test_link_entity_to_investor(auth_client, session):
    inv = make_user(session, username="ten31llc", role=UserRole.investor, name="Ten31 LLC")
    gp = Entity(name="Ten31 LLC", type=EntityType.gp)
    session.add(gp)
    session.commit()
    session.refresh(gp)

    linked = auth_client.patch(f"/api/entities/{gp.id}", json={"linked_user_id": inv.id})
    assert linked.status_code == 200, linked.text
    assert linked.json()["linked_user_id"] == inv.id

    # Unlink.
    unlinked = auth_client.patch(f"/api/entities/{gp.id}", json={"linked_user_id": None})
    assert unlinked.status_code == 200
    assert unlinked.json()["linked_user_id"] is None


def test_link_rejects_non_investor(auth_client, session):
    staff = make_user(session, username="ops2", role=UserRole.operations)
    gp = Entity(name="Mgmt Co", type=EntityType.mgmt_co)
    session.add(gp)
    session.commit()
    session.refresh(gp)
    resp = auth_client.patch(f"/api/entities/{gp.id}", json={"linked_user_id": staff.id})
    assert resp.status_code == 400