fix: fix delete_account_task not check billing enabled (#29577)

This commit is contained in:
wangxiaolei 2025-12-15 10:12:35 +08:00 committed by GitHub
parent 59137f1d05
commit 470650d1d7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 115 additions and 1 deletions

View File

@ -2,6 +2,7 @@ import logging
from celery import shared_task from celery import shared_task
from configs import dify_config
from extensions.ext_database import db from extensions.ext_database import db
from models import Account from models import Account
from services.billing_service import BillingService from services.billing_service import BillingService
@ -14,7 +15,8 @@ logger = logging.getLogger(__name__)
def delete_account_task(account_id): def delete_account_task(account_id):
account = db.session.query(Account).where(Account.id == account_id).first() account = db.session.query(Account).where(Account.id == account_id).first()
try: try:
BillingService.delete_account(account_id) if dify_config.BILLING_ENABLED:
BillingService.delete_account(account_id)
except Exception: except Exception:
logger.exception("Failed to delete account %s from billing service.", account_id) logger.exception("Failed to delete account %s from billing service.", account_id)
raise raise

View File

@ -0,0 +1,112 @@
"""
Unit tests for delete_account_task.
Covers:
- Billing enabled with existing account: calls billing and sends success email
- Billing disabled with existing account: skips billing, sends success email
- Account not found: still calls billing when enabled, does not send email
- Billing deletion raises: logs and re-raises, no email
"""
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from tasks.delete_account_task import delete_account_task
@pytest.fixture
def mock_db_session():
"""Mock the db.session used in delete_account_task."""
with patch("tasks.delete_account_task.db.session") as mock_session:
mock_query = MagicMock()
mock_session.query.return_value = mock_query
mock_query.where.return_value = mock_query
yield mock_session
@pytest.fixture
def mock_deps():
"""Patch external dependencies: BillingService and send_deletion_success_task."""
with (
patch("tasks.delete_account_task.BillingService") as mock_billing,
patch("tasks.delete_account_task.send_deletion_success_task") as mock_mail_task,
):
# ensure .delay exists on the mail task
mock_mail_task.delay = MagicMock()
yield {
"billing": mock_billing,
"mail_task": mock_mail_task,
}
def _set_account_found(mock_db_session, email: str = "user@example.com"):
account = SimpleNamespace(email=email)
mock_db_session.query.return_value.where.return_value.first.return_value = account
return account
def _set_account_missing(mock_db_session):
mock_db_session.query.return_value.where.return_value.first.return_value = None
class TestDeleteAccountTask:
def test_billing_enabled_account_exists_calls_billing_and_sends_email(self, mock_db_session, mock_deps):
# Arrange
account_id = "acc-123"
account = _set_account_found(mock_db_session, email="a@b.com")
# Enable billing
with patch("tasks.delete_account_task.dify_config.BILLING_ENABLED", True):
# Act
delete_account_task(account_id)
# Assert
mock_deps["billing"].delete_account.assert_called_once_with(account_id)
mock_deps["mail_task"].delay.assert_called_once_with(account.email)
def test_billing_disabled_account_exists_sends_email_only(self, mock_db_session, mock_deps):
# Arrange
account_id = "acc-456"
account = _set_account_found(mock_db_session, email="x@y.com")
# Disable billing
with patch("tasks.delete_account_task.dify_config.BILLING_ENABLED", False):
# Act
delete_account_task(account_id)
# Assert
mock_deps["billing"].delete_account.assert_not_called()
mock_deps["mail_task"].delay.assert_called_once_with(account.email)
def test_account_not_found_billing_enabled_calls_billing_no_email(self, mock_db_session, mock_deps, caplog):
# Arrange
account_id = "missing-id"
_set_account_missing(mock_db_session)
# Enable billing
with patch("tasks.delete_account_task.dify_config.BILLING_ENABLED", True):
# Act
delete_account_task(account_id)
# Assert
mock_deps["billing"].delete_account.assert_called_once_with(account_id)
mock_deps["mail_task"].delay.assert_not_called()
# Optional: verify log contains not found message
assert any("not found" in rec.getMessage().lower() for rec in caplog.records)
def test_billing_delete_raises_propagates_and_no_email(self, mock_db_session, mock_deps):
# Arrange
account_id = "acc-err"
_set_account_found(mock_db_session, email="err@ex.com")
mock_deps["billing"].delete_account.side_effect = RuntimeError("billing down")
# Enable billing
with patch("tasks.delete_account_task.dify_config.BILLING_ENABLED", True):
# Act & Assert
with pytest.raises(RuntimeError):
delete_account_task(account_id)
# Ensure email was not sent
mock_deps["mail_task"].delay.assert_not_called()