dify/api/tests/unit_tests/controllers/console/billing/test_billing.py

import base64
import json
from unittest.mock import patch

import pytest
from flask import Flask
from sqlalchemy.orm import Session
from werkzeug.exceptions import BadRequest

from controllers.console import wraps as console_wraps
from controllers.console.billing.billing import PartnerTenants
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.model import DifySetup


@pytest.mark.parametrize(
    "sqlite_session",
    [(DifySetup, Account, Tenant, TenantAccountJoin)],
    indirect=True,
)
class TestPartnerTenants:
    """Unit tests for PartnerTenants controller."""

    @pytest.fixture
    def app(self):
        """Create Flask app for testing."""
        app = Flask(__name__)
        app.config["TESTING"] = True
        app.config["SECRET_KEY"] = "test-secret-key"
        return app

    @pytest.fixture
    def mock_account(self, sqlite_session: Session):
        """Persist an initialized account with an owner workspace membership."""
        tenant = Tenant(name="Billing Tenant")
        account = Account(name="Billing User", email="test@example.com")
        sqlite_session.add_all([tenant, account])
        sqlite_session.flush()
        sqlite_session.add_all(
            [
                TenantAccountJoin(
                    tenant_id=tenant.id,
                    account_id=account.id,
                    current=True,
                    role=TenantAccountRole.OWNER,
                    invited_by=None,
                ),
                DifySetup(version="test"),
            ]
        )
        sqlite_session.commit()
        account._current_tenant = tenant
        sqlite_session.expunge(account)
        return account

    @pytest.fixture
    def mock_billing_service(self):
        """Mock BillingService."""
        with patch("controllers.console.billing.billing.BillingService") as mock_service:
            yield mock_service

    @pytest.fixture
    def mock_decorators(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session):
        """Keep authentication mocked while the setup guard uses SQLite."""
        console_wraps._is_setup_completed.reset_success()
        monkeypatch.setattr(console_wraps.db, "session", sqlite_session)
        with (
            patch("controllers.console.wraps.dify_config.EDITION", "CLOUD"),
            patch("libs.login.dify_config.LOGIN_DISABLED", False),
            patch("libs.login.check_csrf_token") as mock_csrf,
        ):
            mock_csrf.return_value = None
            yield mock_csrf
        console_wraps._is_setup_completed.reset_success()

    def test_put_success(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
        """Test successful partner tenants bindings sync."""
        # Arrange
        partner_key_encoded = base64.b64encode(b"partner-key-123").decode("utf-8")
        click_id = "click-id-789"
        expected_response = {"result": "success", "data": {"synced": True}}

        mock_billing_service.sync_partner_tenants_bindings.return_value = expected_response

        with app.test_request_context(
            method="PUT",
            json={"click_id": click_id},
            path=f"/billing/partners/{partner_key_encoded}/tenants",
        ):
            with (
                patch(
                    "controllers.console.wraps.current_account_with_tenant",
                    return_value=(mock_account, mock_account.current_tenant_id),
                ),
                patch("libs.login._get_user", return_value=mock_account),
            ):
                resource = PartnerTenants()
                result = resource.put(partner_key_encoded)

        # Assert
        assert result == expected_response
        mock_billing_service.sync_partner_tenants_bindings.assert_called_once_with(
            mock_account.id, "partner-key-123", click_id
        )

    def test_put_invalid_partner_key_base64(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
        """Test that invalid base64 partner_key raises BadRequest."""
        # Arrange
        invalid_partner_key = "invalid-base64-!@#$"
        click_id = "click-id-789"

        with app.test_request_context(
            method="PUT",
            json={"click_id": click_id},
            path=f"/billing/partners/{invalid_partner_key}/tenants",
        ):
            with (
                patch(
                    "controllers.console.wraps.current_account_with_tenant",
                    return_value=(mock_account, mock_account.current_tenant_id),
                ),
                patch("libs.login._get_user", return_value=mock_account),
            ):
                resource = PartnerTenants()

                # Act & Assert
                with pytest.raises(BadRequest) as exc_info:
                    resource.put(invalid_partner_key)
                assert "Invalid partner_key" in str(exc_info.value)

    def test_put_missing_click_id(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
        """Test that missing click_id raises BadRequest."""
        # Arrange
        partner_key_encoded = base64.b64encode(b"partner-key-123").decode("utf-8")

        with app.test_request_context(
            method="PUT",
            json={},
            path=f"/billing/partners/{partner_key_encoded}/tenants",
        ):
            with (
                patch(
                    "controllers.console.wraps.current_account_with_tenant",
                    return_value=(mock_account, mock_account.current_tenant_id),
                ),
                patch("libs.login._get_user", return_value=mock_account),
            ):
                resource = PartnerTenants()

                # Act & Assert
                # Validation should raise BadRequest for missing required field
                with pytest.raises(BadRequest):
                    resource.put(partner_key_encoded)

    def test_put_billing_service_json_decode_error(
        self, app: Flask, mock_account, mock_billing_service, mock_decorators
    ):
        """Test handling of billing service JSON decode error.

        When billing service returns non-200 status code with invalid JSON response,
        response.json() raises JSONDecodeError. This exception propagates to the controller
        and should be handled by the global error handler (handle_general_exception),
        which returns a 500 status code with error details.

        Note: In unit tests, when directly calling resource.put(), the exception is raised
        directly. In actual Flask application, the error handler would catch it and return
        a 500 response with JSON: {"code": "unknown", "message": "...", "status": 500}
        """
        # Arrange
        partner_key_encoded = base64.b64encode(b"partner-key-123").decode("utf-8")
        click_id = "click-id-789"

        # Simulate JSON decode error when billing service returns invalid JSON
        # This happens when billing service returns non-200 with empty/invalid response body
        json_decode_error = json.JSONDecodeError("Expecting value", "", 0)
        mock_billing_service.sync_partner_tenants_bindings.side_effect = json_decode_error

        with app.test_request_context(
            method="PUT",
            json={"click_id": click_id},
            path=f"/billing/partners/{partner_key_encoded}/tenants",
        ):
            with (
                patch(
                    "controllers.console.wraps.current_account_with_tenant",
                    return_value=(mock_account, mock_account.current_tenant_id),
                ),
                patch("libs.login._get_user", return_value=mock_account),
            ):
                resource = PartnerTenants()

                # Act & Assert
                # JSONDecodeError will be raised from the controller
                # In actual Flask app, this would be caught by handle_general_exception
                # which returns: {"code": "unknown", "message": str(e), "status": 500}
                with pytest.raises(json.JSONDecodeError) as exc_info:
                    resource.put(partner_key_encoded)

                # Verify the exception is JSONDecodeError
                assert isinstance(exc_info.value, json.JSONDecodeError)
                assert "Expecting value" in str(exc_info.value)

    def test_put_empty_click_id(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
        """Test that empty click_id raises BadRequest."""
        # Arrange
        partner_key_encoded = base64.b64encode(b"partner-key-123").decode("utf-8")
        click_id = ""

        with app.test_request_context(
            method="PUT",
            json={"click_id": click_id},
            path=f"/billing/partners/{partner_key_encoded}/tenants",
        ):
            with (
                patch(
                    "controllers.console.wraps.current_account_with_tenant",
                    return_value=(mock_account, mock_account.current_tenant_id),
                ),
                patch("libs.login._get_user", return_value=mock_account),
            ):
                resource = PartnerTenants()

                # Act & Assert
                with pytest.raises(BadRequest) as exc_info:
                    resource.put(partner_key_encoded)
                assert "Invalid partner information" in str(exc_info.value)

    def test_put_empty_partner_key_after_decode(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
        """Test that empty partner_key after decode raises BadRequest."""
        # Arrange
        # Base64 encode an empty string
        empty_partner_key_encoded = base64.b64encode(b"").decode("utf-8")
        click_id = "click-id-789"

        with app.test_request_context(
            method="PUT",
            json={"click_id": click_id},
            path=f"/billing/partners/{empty_partner_key_encoded}/tenants",
        ):
            with (
                patch(
                    "controllers.console.wraps.current_account_with_tenant",
                    return_value=(mock_account, mock_account.current_tenant_id),
                ),
                patch("libs.login._get_user", return_value=mock_account),
            ):
                resource = PartnerTenants()

                # Act & Assert
                with pytest.raises(BadRequest) as exc_info:
                    resource.put(empty_partner_key_encoded)
                assert "Invalid partner information" in str(exc_info.value)

    def test_put_empty_user_id(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
        """Test that empty user id raises BadRequest."""
        # Arrange
        partner_key_encoded = base64.b64encode(b"partner-key-123").decode("utf-8")
        click_id = "click-id-789"
        mock_account.id = None  # Empty user id

        with app.test_request_context(
            method="PUT",
            json={"click_id": click_id},
            path=f"/billing/partners/{partner_key_encoded}/tenants",
        ):
            with (
                patch(
                    "controllers.console.wraps.current_account_with_tenant",
                    return_value=(mock_account, mock_account.current_tenant_id),
                ),
                patch("libs.login._get_user", return_value=mock_account),
            ):
                resource = PartnerTenants()

                # Act & Assert
                with pytest.raises(BadRequest) as exc_info:
                    resource.put(partner_key_encoded)
                assert "Invalid partner information" in str(exc_info.value)