Feat integrate partner stack (#28353)
Co-authored-by: Joel <iamjoel007@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: crazywoola <100913391+crazywoola@users.noreply.github.com>
This commit is contained in:
hj24
253
api/tests/unit_tests/controllers/console/billing/test_billing.py
Normal file
253
api/tests/unit_tests/controllers/console/billing/test_billing.py
Normal file
@@ -0,0 +1,253 @@
|
||||
import base64
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from werkzeug.exceptions import BadRequest
|
||||
|
||||
from controllers.console.billing.billing import PartnerTenants
|
||||
from models.account import Account
|
||||
|
||||
|
||||
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):
|
||||
"""Create a mock account."""
|
||||
account = MagicMock(spec=Account)
|
||||
account.id = "account-123"
|
||||
account.email = "test@example.com"
|
||||
account.current_tenant_id = "tenant-456"
|
||||
account.is_authenticated = True
|
||||
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):
|
||||
"""Mock decorators to avoid database access."""
|
||||
with (
|
||||
patch("controllers.console.wraps.db") as mock_db,
|
||||
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_db.session.query.return_value.first.return_value = MagicMock() # Mock setup exists
|
||||
mock_csrf.return_value = None
|
||||
yield {"db": mock_db, "csrf": mock_csrf}
|
||||
|
||||
def test_put_success(self, app, 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.billing.billing.current_account_with_tenant",
|
||||
return_value=(mock_account, "tenant-456"),
|
||||
),
|
||||
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, 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.billing.billing.current_account_with_tenant",
|
||||
return_value=(mock_account, "tenant-456"),
|
||||
),
|
||||
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, 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.billing.billing.current_account_with_tenant",
|
||||
return_value=(mock_account, "tenant-456"),
|
||||
),
|
||||
patch("libs.login._get_user", return_value=mock_account),
|
||||
):
|
||||
resource = PartnerTenants()
|
||||
|
||||
# Act & Assert
|
||||
# reqparse will 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, 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.billing.billing.current_account_with_tenant",
|
||||
return_value=(mock_account, "tenant-456"),
|
||||
),
|
||||
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, 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.billing.billing.current_account_with_tenant",
|
||||
return_value=(mock_account, "tenant-456"),
|
||||
),
|
||||
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, 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.billing.billing.current_account_with_tenant",
|
||||
return_value=(mock_account, "tenant-456"),
|
||||
),
|
||||
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, 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.billing.billing.current_account_with_tenant",
|
||||
return_value=(mock_account, "tenant-456"),
|
||||
),
|
||||
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)
|
||||
236
api/tests/unit_tests/services/test_billing_service.py
Normal file
236
api/tests/unit_tests/services/test_billing_service.py
Normal file
@@ -0,0 +1,236 @@
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from werkzeug.exceptions import InternalServerError
|
||||
|
||||
from services.billing_service import BillingService
|
||||
|
||||
|
||||
class TestBillingServiceSendRequest:
|
||||
"""Unit tests for BillingService._send_request method."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_httpx_request(self):
|
||||
"""Mock httpx.request for testing."""
|
||||
with patch("services.billing_service.httpx.request") as mock_request:
|
||||
yield mock_request
|
||||
|
||||
@pytest.fixture
|
||||
def mock_billing_config(self):
|
||||
"""Mock BillingService configuration."""
|
||||
with (
|
||||
patch.object(BillingService, "base_url", "https://billing-api.example.com"),
|
||||
patch.object(BillingService, "secret_key", "test-secret-key"),
|
||||
):
|
||||
yield
|
||||
|
||||
def test_get_request_success(self, mock_httpx_request, mock_billing_config):
|
||||
"""Test successful GET request."""
|
||||
# Arrange
|
||||
expected_response = {"result": "success", "data": {"info": "test"}}
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = httpx.codes.OK
|
||||
mock_response.json.return_value = expected_response
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act
|
||||
result = BillingService._send_request("GET", "/test", params={"key": "value"})
|
||||
|
||||
# Assert
|
||||
assert result == expected_response
|
||||
mock_httpx_request.assert_called_once()
|
||||
call_args = mock_httpx_request.call_args
|
||||
assert call_args[0][0] == "GET"
|
||||
assert call_args[0][1] == "https://billing-api.example.com/test"
|
||||
assert call_args[1]["params"] == {"key": "value"}
|
||||
assert call_args[1]["headers"]["Billing-Api-Secret-Key"] == "test-secret-key"
|
||||
assert call_args[1]["headers"]["Content-Type"] == "application/json"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code", [httpx.codes.NOT_FOUND, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.BAD_REQUEST]
|
||||
)
|
||||
def test_get_request_non_200_status_code(self, mock_httpx_request, mock_billing_config, status_code):
|
||||
"""Test GET request with non-200 status code raises ValueError."""
|
||||
# Arrange
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = status_code
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
BillingService._send_request("GET", "/test")
|
||||
assert "Unable to retrieve billing information" in str(exc_info.value)
|
||||
|
||||
def test_put_request_success(self, mock_httpx_request, mock_billing_config):
|
||||
"""Test successful PUT request."""
|
||||
# Arrange
|
||||
expected_response = {"result": "success"}
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = httpx.codes.OK
|
||||
mock_response.json.return_value = expected_response
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act
|
||||
result = BillingService._send_request("PUT", "/test", json={"key": "value"})
|
||||
|
||||
# Assert
|
||||
assert result == expected_response
|
||||
call_args = mock_httpx_request.call_args
|
||||
assert call_args[0][0] == "PUT"
|
||||
|
||||
def test_put_request_internal_server_error(self, mock_httpx_request, mock_billing_config):
|
||||
"""Test PUT request with INTERNAL_SERVER_ERROR raises InternalServerError."""
|
||||
# Arrange
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = httpx.codes.INTERNAL_SERVER_ERROR
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(InternalServerError) as exc_info:
|
||||
BillingService._send_request("PUT", "/test", json={"key": "value"})
|
||||
assert exc_info.value.code == 500
|
||||
assert "Unable to process billing request" in str(exc_info.value.description)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code", [httpx.codes.BAD_REQUEST, httpx.codes.NOT_FOUND, httpx.codes.UNAUTHORIZED, httpx.codes.FORBIDDEN]
|
||||
)
|
||||
def test_put_request_non_200_non_500(self, mock_httpx_request, mock_billing_config, status_code):
|
||||
"""Test PUT request with non-200 and non-500 status code raises ValueError."""
|
||||
# Arrange
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = status_code
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
BillingService._send_request("PUT", "/test", json={"key": "value"})
|
||||
assert "Invalid arguments." in str(exc_info.value)
|
||||
|
||||
@pytest.mark.parametrize("method", ["POST", "DELETE"])
|
||||
def test_non_get_non_put_request_success(self, mock_httpx_request, mock_billing_config, method):
|
||||
"""Test successful POST/DELETE request."""
|
||||
# Arrange
|
||||
expected_response = {"result": "success"}
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = httpx.codes.OK
|
||||
mock_response.json.return_value = expected_response
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act
|
||||
result = BillingService._send_request(method, "/test", json={"key": "value"})
|
||||
|
||||
# Assert
|
||||
assert result == expected_response
|
||||
call_args = mock_httpx_request.call_args
|
||||
assert call_args[0][0] == method
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code", [httpx.codes.BAD_REQUEST, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.NOT_FOUND]
|
||||
)
|
||||
def test_post_request_non_200_with_valid_json(self, mock_httpx_request, mock_billing_config, status_code):
|
||||
"""Test POST request with non-200 status code raises ValueError."""
|
||||
# Arrange
|
||||
error_response = {"detail": "Error message"}
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = status_code
|
||||
mock_response.json.return_value = error_response
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
BillingService._send_request("POST", "/test", json={"key": "value"})
|
||||
assert "Unable to send request to" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code", [httpx.codes.BAD_REQUEST, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.NOT_FOUND]
|
||||
)
|
||||
def test_delete_request_non_200_with_valid_json(self, mock_httpx_request, mock_billing_config, status_code):
|
||||
"""Test DELETE request with non-200 status code but valid JSON response.
|
||||
|
||||
DELETE doesn't check status code, so it returns the error JSON.
|
||||
"""
|
||||
# Arrange
|
||||
error_response = {"detail": "Error message"}
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = status_code
|
||||
mock_response.json.return_value = error_response
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act
|
||||
result = BillingService._send_request("DELETE", "/test", json={"key": "value"})
|
||||
|
||||
# Assert
|
||||
assert result == error_response
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code", [httpx.codes.BAD_REQUEST, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.NOT_FOUND]
|
||||
)
|
||||
def test_post_request_non_200_with_invalid_json(self, mock_httpx_request, mock_billing_config, status_code):
|
||||
"""Test POST request with non-200 status code raises ValueError before JSON parsing."""
|
||||
# Arrange
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = status_code
|
||||
mock_response.text = ""
|
||||
mock_response.json.side_effect = json.JSONDecodeError("Expecting value", "", 0)
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act & Assert
|
||||
# POST checks status code before calling response.json(), so ValueError is raised
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
BillingService._send_request("POST", "/test", json={"key": "value"})
|
||||
assert "Unable to send request to" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code", [httpx.codes.BAD_REQUEST, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.NOT_FOUND]
|
||||
)
|
||||
def test_delete_request_non_200_with_invalid_json(self, mock_httpx_request, mock_billing_config, status_code):
|
||||
"""Test DELETE request with non-200 status code and invalid JSON response raises exception.
|
||||
|
||||
DELETE doesn't check status code, so it calls response.json() which raises JSONDecodeError
|
||||
when the response cannot be parsed as JSON (e.g., empty response).
|
||||
"""
|
||||
# Arrange
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = status_code
|
||||
mock_response.text = ""
|
||||
mock_response.json.side_effect = json.JSONDecodeError("Expecting value", "", 0)
|
||||
mock_httpx_request.return_value = mock_response
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
BillingService._send_request("DELETE", "/test", json={"key": "value"})
|
||||
|
||||
def test_retry_on_request_error(self, mock_httpx_request, mock_billing_config):
|
||||
"""Test that _send_request retries on httpx.RequestError."""
|
||||
# Arrange
|
||||
expected_response = {"result": "success"}
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = httpx.codes.OK
|
||||
mock_response.json.return_value = expected_response
|
||||
|
||||
# First call raises RequestError, second succeeds
|
||||
mock_httpx_request.side_effect = [
|
||||
httpx.RequestError("Network error"),
|
||||
mock_response,
|
||||
]
|
||||
|
||||
# Act
|
||||
result = BillingService._send_request("GET", "/test")
|
||||
|
||||
# Assert
|
||||
assert result == expected_response
|
||||
assert mock_httpx_request.call_count == 2
|
||||
|
||||
def test_retry_exhausted_raises_exception(self, mock_httpx_request, mock_billing_config):
|
||||
"""Test that _send_request raises exception after retries are exhausted."""
|
||||
# Arrange
|
||||
mock_httpx_request.side_effect = httpx.RequestError("Network error")
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(httpx.RequestError):
|
||||
BillingService._send_request("GET", "/test")
|
||||
|
||||
# Should retry multiple times (wait=2, stop_before_delay=10 means ~5 attempts)
|
||||
assert mock_httpx_request.call_count > 1
|
||||
Reference in New Issue
Block a user