admin/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat: trigger billing (#28335)

Browse Source 
Signed-off-by: lyzno1 <yuanyouhuilyz@gmail.com>
Co-authored-by: lyzno1 <yuanyouhuilyz@gmail.com>
Co-authored-by: lyzno1 <92089059+lyzno1@users.noreply.github.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Maries
2025-11-20 10:15:23 +08:00
committed by GitHub
parent c0b7ffd5d0
commit a1b735a4c0
61 changed files with 1475 additions and 465 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
api/tests/test_containers_integration_tests/services/test_app_generate_service.py
View File

@@ -5,12 +5,10 @@ import pytest
from faker import Faker
from core.app.entities.app_invoke_entities import InvokeFrom
from enums.cloud_plan import CloudPlan
from models.model import EndUser
from models.workflow import Workflow
from services.app_generate_service import AppGenerateService
from services.errors.app import WorkflowIdFormatError, WorkflowNotFoundError
from services.errors.llm import InvokeRateLimitError
class TestAppGenerateService:
@@ -20,10 +18,9 @@ class TestAppGenerateService:
def mock_external_service_dependencies(self):
"""Mock setup for external service dependencies."""
with (
patch("services.app_generate_service.BillingService") as mock_billing_service,
patch("services.billing_service.BillingService") as mock_billing_service,
patch("services.app_generate_service.WorkflowService") as mock_workflow_service,
patch("services.app_generate_service.RateLimit") as mock_rate_limit,
patch("services.app_generate_service.RateLimiter") as mock_rate_limiter,
patch("services.app_generate_service.CompletionAppGenerator") as mock_completion_generator,
patch("services.app_generate_service.ChatAppGenerator") as mock_chat_generator,
patch("services.app_generate_service.AgentChatAppGenerator") as mock_agent_chat_generator,
@@ -31,9 +28,13 @@ class TestAppGenerateService:
patch("services.app_generate_service.WorkflowAppGenerator") as mock_workflow_generator,
patch("services.account_service.FeatureService") as mock_account_feature_service,
patch("services.app_generate_service.dify_config") as mock_dify_config,
patch("configs.dify_config") as mock_global_dify_config,
):
# Setup default mock returns for billing service
mock_billing_service.get_info.return_value = {"subscription": {"plan": CloudPlan.SANDBOX}}
mock_billing_service.update_tenant_feature_plan_usage.return_value = {
"result": "success",
"history_id": "test_history_id",
}
# Setup default mock returns for workflow service
mock_workflow_service_instance = mock_workflow_service.return_value
@@ -47,10 +48,6 @@ class TestAppGenerateService:
mock_rate_limit_instance.generate.return_value = ["test_response"]
mock_rate_limit_instance.exit.return_value = None
mock_rate_limiter_instance = mock_rate_limiter.return_value
mock_rate_limiter_instance.is_rate_limited.return_value = False
mock_rate_limiter_instance.increment_rate_limit.return_value = None
# Setup default mock returns for app generators
mock_completion_generator_instance = mock_completion_generator.return_value
mock_completion_generator_instance.generate.return_value = ["completion_response"]
@@ -87,11 +84,14 @@ class TestAppGenerateService:
mock_dify_config.APP_MAX_ACTIVE_REQUESTS = 100
mock_dify_config.APP_DAILY_RATE_LIMIT = 1000
mock_global_dify_config.BILLING_ENABLED = False
mock_global_dify_config.APP_MAX_ACTIVE_REQUESTS = 100
mock_global_dify_config.APP_DAILY_RATE_LIMIT = 1000
yield {
"billing_service": mock_billing_service,
"workflow_service": mock_workflow_service,
"rate_limit": mock_rate_limit,
"rate_limiter": mock_rate_limiter,
"completion_generator": mock_completion_generator,
"chat_generator": mock_chat_generator,
"agent_chat_generator": mock_agent_chat_generator,
@@ -99,6 +99,7 @@ class TestAppGenerateService:
"workflow_generator": mock_workflow_generator,
"account_feature_service": mock_account_feature_service,
"dify_config": mock_dify_config,
"global_dify_config": mock_global_dify_config,
}
def _create_test_app_and_account(self, db_session_with_containers, mock_external_service_dependencies, mode="chat"):
@@ -429,13 +430,9 @@ class TestAppGenerateService:
db_session_with_containers, mock_external_service_dependencies, mode="completion"
)
# Setup billing service mock for sandbox plan
mock_external_service_dependencies["billing_service"].get_info.return_value = {
"subscription": {"plan": CloudPlan.SANDBOX}
}
# Set BILLING_ENABLED to True for this test
mock_external_service_dependencies["dify_config"].BILLING_ENABLED = True
mock_external_service_dependencies["global_dify_config"].BILLING_ENABLED = True
# Setup test arguments
args = {"inputs": {"query": fake.text(max_nb_chars=50)}, "response_mode": "streaming"}
@@ -448,41 +445,8 @@ class TestAppGenerateService:
# Verify the result
assert result == ["test_response"]
# Verify billing service was called
mock_external_service_dependencies["billing_service"].get_info.assert_called_once_with(app.tenant_id)
def test_generate_with_rate_limit_exceeded(self, db_session_with_containers, mock_external_service_dependencies):
"""
Test generation when rate limit is exceeded.
"""
fake = Faker()
app, account = self._create_test_app_and_account(
db_session_with_containers, mock_external_service_dependencies, mode="completion"
)
# Setup billing service mock for sandbox plan
mock_external_service_dependencies["billing_service"].get_info.return_value = {
"subscription": {"plan": CloudPlan.SANDBOX}
}
# Set BILLING_ENABLED to True for this test
mock_external_service_dependencies["dify_config"].BILLING_ENABLED = True
# Setup system rate limiter to return rate limited
with patch("services.app_generate_service.AppGenerateService.system_rate_limiter") as mock_system_rate_limiter:
mock_system_rate_limiter.is_rate_limited.return_value = True
# Setup test arguments
args = {"inputs": {"query": fake.text(max_nb_chars=50)}, "response_mode": "streaming"}
# Execute the method under test and expect rate limit error
with pytest.raises(InvokeRateLimitError) as exc_info:
AppGenerateService.generate(
app_model=app, user=account, args=args, invoke_from=InvokeFrom.SERVICE_API, streaming=True
)
# Verify error message
assert "Rate limit exceeded" in str(exc_info.value)
# Verify billing service was called to consume quota
mock_external_service_dependencies["billing_service"].update_tenant_feature_plan_usage.assert_called_once()
def test_generate_with_invalid_app_mode(self, db_session_with_containers, mock_external_service_dependencies):
"""