"""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 that free plan has tree creation limit of 3.""" tree_template = { "name": "Limit Test Tree", "tree_structure": { "id": "root", "type": "solution", "title": "Test", "description": "Test tree" } } # 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" } } 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" }, "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 that free plan limits are correct.""" 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