resolutionflow/backend/tests/test_subscription_limits.py

"""Integration tests for subscription limits."""

import pytest
from httpx import AsyncClient
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession


class TestSubscriptionLimits:
    """Test suite for subscription plan limits."""

    @pytest.mark.asyncio
    async def test_free_plan_tree_limit(
        self, client: AsyncClient, auth_headers: dict, test_db: AsyncSession
    ):
        """Test that free plan has tree creation limit of 3."""
        from app.models.subscription import Subscription
        sub = (await test_db.execute(select(Subscription))).scalar_one()
        sub.plan = "free"
        await test_db.commit()

        tree_template = {
            "name": "Limit Test Tree",
            "tree_structure": {
                "id": "root",
                "type": "solution",
                "title": "Test",
                "description": "Test tree",
                "solution": "Test solution"
            }
        }

        # Create trees up to the limit
        for i in range(3):
            tree_data = {**tree_template, "name": f"Tree {i+1}"}
            response = await client.post("/api/v1/trees", json=tree_data, headers=auth_headers)
            assert response.status_code == 201, f"Failed creating tree {i+1}: {response.json()}"

        # 4th tree should fail with 402
        tree_data = {**tree_template, "name": "Tree 4 Over Limit"}
        response = await client.post("/api/v1/trees", json=tree_data, headers=auth_headers)
        assert response.status_code == 402
        assert "limit" in response.json()["detail"].lower()

    @pytest.mark.asyncio
    async def test_subscription_details_show_usage(self, client: AsyncClient, auth_headers: dict):
        """Test that subscription details reflect actual usage."""
        # Check initial usage
        response = await client.get("/api/v1/accounts/me/subscription", headers=auth_headers)
        assert response.status_code == 200
        initial_usage = response.json()["usage"]
        assert initial_usage["tree_count"] == 0

        # Create a tree
        tree_data = {
            "name": "Usage Test Tree",
            "tree_structure": {
                "id": "root",
                "type": "solution",
                "title": "Test",
                "description": "Test",
                "solution": "Test solution"
            }
        }
        create_resp = await client.post("/api/v1/trees", json=tree_data, headers=auth_headers)
        assert create_resp.status_code == 201

        # Check usage increased
        response = await client.get("/api/v1/accounts/me/subscription", headers=auth_headers)
        assert response.status_code == 200
        updated_usage = response.json()["usage"]
        assert updated_usage["tree_count"] == 1

    @pytest.mark.asyncio
    async def test_super_admin_bypasses_limits(
        self, client: AsyncClient, admin_auth_headers: dict
    ):
        """Test that super admin can create trees without limit checks."""
        tree_template = {
            "name": "Admin Tree",
            "tree_structure": {
                "id": "root",
                "type": "solution",
                "title": "Test",
                "description": "Test tree",
                "solution": "Test solution"
            },
            "is_default": True  # Default trees skip limit check
        }

        # Super admin creating default trees should always work
        for i in range(5):
            tree_data = {**tree_template, "name": f"Admin Tree {i+1}"}
            response = await client.post(
                "/api/v1/trees", json=tree_data, headers=admin_auth_headers
            )
            assert response.status_code == 201

    @pytest.mark.asyncio
    async def test_free_plan_limits_correct(
        self, client: AsyncClient, auth_headers: dict, test_db: AsyncSession
    ):
        """Test that free plan limits are correct."""
        from app.models.subscription import Subscription
        sub = (await test_db.execute(select(Subscription))).scalar_one()
        sub.plan = "free"
        await test_db.commit()

        response = await client.get("/api/v1/accounts/me/subscription", headers=auth_headers)
        assert response.status_code == 200
        limits = response.json()["limits"]
        assert limits["plan"] == "free"
        assert limits["max_trees"] == 3
        assert limits["max_sessions_per_month"] == 20
        assert limits["max_users"] == 1
        assert limits["custom_branding"] is False
        assert limits["priority_support"] is False

    @pytest.mark.asyncio
    async def test_upgraded_plan_has_higher_limits(
        self, client: AsyncClient, auth_headers: dict, test_db: AsyncSession
    ):
        """Test that upgrading plan increases limits."""
        from app.models.subscription import Subscription
        from app.models.user import User

        # Get the user's subscription and upgrade it
        me_resp = await client.get("/api/v1/auth/me", headers=auth_headers)
        account_id_str = me_resp.json()["account_id"]

        from uuid import UUID
        account_id = UUID(account_id_str)
        result = await test_db.execute(
            select(Subscription).where(Subscription.account_id == account_id)
        )
        sub = result.scalar_one()
        sub.plan = "pro"
        await test_db.commit()

        # Check limits are now pro
        response = await client.get("/api/v1/accounts/me/subscription", headers=auth_headers)
        assert response.status_code == 200
        limits = response.json()["limits"]
        assert limits["plan"] == "pro"
        assert limits["max_trees"] == 25
        assert limits["max_sessions_per_month"] == 200