""" Unit tests for Service API wraps (authentication decorators) """ import uuid from unittest.mock import MagicMock, Mock, patch import pytest from flask import Flask from werkzeug.exceptions import Forbidden, NotFound, Unauthorized from controllers.service_api.wraps import ( DatasetApiResource, FetchUserArg, WhereisUserArg, cloud_edition_billing_knowledge_limit_check, cloud_edition_billing_rate_limit_check, cloud_edition_billing_resource_check, validate_and_get_api_token, validate_app_token, validate_dataset_token, ) from enums.cloud_plan import CloudPlan from models.account import TenantStatus from models.model import ApiToken from tests.unit_tests.conftest import ( setup_mock_dataset_owner_execute_result, setup_mock_tenant_owner_execute_result, ) def _configure_current_app_mock(mock_current_app): mock_current_app.login_manager = Mock() mock_current_app._get_current_object = Mock(return_value=Mock()) class TestValidateAndGetApiToken: """Test suite for validate_and_get_api_token function""" @pytest.fixture def app(self): """Create Flask test application.""" app = Flask(__name__) app.config["TESTING"] = True return app def test_missing_authorization_header(self, app: Flask): """Test that Unauthorized is raised when Authorization header is missing.""" # Arrange with app.test_request_context("/", method="GET"): # No Authorization header # Act & Assert with pytest.raises(Unauthorized) as exc_info: validate_and_get_api_token("app") assert "Authorization header must be provided" in str(exc_info.value) def test_invalid_auth_scheme(self, app: Flask): """Test that Unauthorized is raised when auth scheme is not Bearer.""" # Arrange with app.test_request_context("/", method="GET", headers={"Authorization": "Basic token123"}): # Act & Assert with pytest.raises(Unauthorized) as exc_info: validate_and_get_api_token("app") assert "Authorization scheme must be 'Bearer'" in str(exc_info.value) @patch("controllers.service_api.wraps.record_token_usage") @patch("controllers.service_api.wraps.ApiTokenCache") @patch("controllers.service_api.wraps.fetch_token_with_single_flight") def test_valid_token_returns_api_token(self, mock_fetch_token, mock_cache_cls, mock_record_usage, app: Flask): """Test that valid token returns the ApiToken object.""" # Arrange mock_api_token = Mock(spec=ApiToken) mock_api_token.token = "valid_token_123" mock_api_token.type = "app" mock_cache_instance = Mock() mock_cache_instance.get.return_value = None # Cache miss mock_cache_cls.get = mock_cache_instance.get mock_fetch_token.return_value = mock_api_token # Act with app.test_request_context("/", method="GET", headers={"Authorization": "Bearer valid_token_123"}): result = validate_and_get_api_token("app") # Assert assert result == mock_api_token @patch("controllers.service_api.wraps.record_token_usage") @patch("controllers.service_api.wraps.ApiTokenCache") @patch("controllers.service_api.wraps.fetch_token_with_single_flight") def test_invalid_token_raises_unauthorized(self, mock_fetch_token, mock_cache_cls, mock_record_usage, app: Flask): """Test that invalid token raises Unauthorized.""" # Arrange from werkzeug.exceptions import Unauthorized mock_cache_instance = Mock() mock_cache_instance.get.return_value = None # Cache miss mock_cache_cls.get = mock_cache_instance.get mock_fetch_token.side_effect = Unauthorized("Access token is invalid") # Act & Assert with app.test_request_context("/", method="GET", headers={"Authorization": "Bearer invalid_token"}): with pytest.raises(Unauthorized) as exc_info: validate_and_get_api_token("app") assert "Access token is invalid" in str(exc_info.value) class TestValidateAppToken: """Test suite for validate_app_token decorator""" @pytest.fixture def app(self): """Create Flask test application.""" app = Flask(__name__) app.config["TESTING"] = True return app @patch("controllers.service_api.wraps.user_logged_in") @patch("controllers.service_api.wraps.db") @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.current_app") def test_valid_app_token_allows_access( self, mock_current_app, mock_validate_token, mock_db, mock_user_logged_in, app ): """Test that valid app token allows access to decorated view.""" # Arrange _configure_current_app_mock(mock_current_app) mock_api_token = Mock() mock_api_token.app_id = str(uuid.uuid4()) mock_api_token.tenant_id = str(uuid.uuid4()) mock_validate_token.return_value = mock_api_token mock_app = Mock() mock_app.id = mock_api_token.app_id mock_app.status = "normal" mock_app.enable_api = True mock_app.tenant_id = mock_api_token.tenant_id mock_tenant = Mock() mock_tenant.status = TenantStatus.NORMAL mock_tenant.id = mock_api_token.tenant_id mock_account = Mock() mock_account.id = str(uuid.uuid4()) # Use side_effect to return app first, then tenant via session.get() mock_db.session.get.side_effect = [mock_app, mock_tenant] # Mock the tenant owner execute result (execute(select(...)).one_or_none()) setup_mock_tenant_owner_execute_result(mock_db, mock_tenant, mock_account) @validate_app_token def protected_view(app_model): return {"success": True, "app_id": app_model.id} # Act with app.test_request_context("/", method="GET", headers={"Authorization": "Bearer test_token"}): result = protected_view() # Assert assert result["success"] is True assert result["app_id"] == mock_app.id @patch("controllers.service_api.wraps.db") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_app_not_found_raises_forbidden(self, mock_validate_token, mock_db, app: Flask): """Test that Forbidden is raised when app no longer exists.""" # Arrange mock_api_token = Mock() mock_api_token.app_id = str(uuid.uuid4()) mock_validate_token.return_value = mock_api_token mock_db.session.get.return_value = None @validate_app_token def protected_view(**kwargs): return {"success": True} # Act & Assert with app.test_request_context("/", method="GET"): with pytest.raises(Forbidden) as exc_info: protected_view() assert "no longer exists" in str(exc_info.value) @patch("controllers.service_api.wraps.db") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_app_status_abnormal_raises_forbidden(self, mock_validate_token, mock_db, app: Flask): """Test that Forbidden is raised when app status is abnormal.""" # Arrange mock_api_token = Mock() mock_api_token.app_id = str(uuid.uuid4()) mock_validate_token.return_value = mock_api_token mock_app = Mock() mock_app.status = "abnormal" mock_db.session.get.return_value = mock_app @validate_app_token def protected_view(**kwargs): return {"success": True} # Act & Assert with app.test_request_context("/", method="GET"): with pytest.raises(Forbidden) as exc_info: protected_view() assert "status is abnormal" in str(exc_info.value) @patch("controllers.service_api.wraps.db") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_app_api_disabled_raises_forbidden(self, mock_validate_token, mock_db, app: Flask): """Test that Forbidden is raised when app API is disabled.""" # Arrange mock_api_token = Mock() mock_api_token.app_id = str(uuid.uuid4()) mock_validate_token.return_value = mock_api_token mock_app = Mock() mock_app.status = "normal" mock_app.enable_api = False mock_db.session.get.return_value = mock_app @validate_app_token def protected_view(**kwargs): return {"success": True} # Act & Assert with app.test_request_context("/", method="GET"): with pytest.raises(Forbidden) as exc_info: protected_view() assert "API service has been disabled" in str(exc_info.value) class TestCloudEditionBillingResourceCheck: """Test suite for cloud_edition_billing_resource_check decorator""" @pytest.fixture def app(self): """Create Flask test application.""" app = Flask(__name__) app.config["TESTING"] = True return app @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_allows_when_under_limit(self, mock_get_features, mock_validate_token, app: Flask): """Test that request is allowed when under resource limit.""" # Arrange mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() mock_features.billing.enabled = True mock_features.members.limit = 10 mock_features.members.size = 5 mock_get_features.return_value = mock_features @cloud_edition_billing_resource_check("members", "app") def add_member(): return "member_added" # Act with app.test_request_context("/", method="GET"): result = add_member() # Assert assert result == "member_added" mock_get_features.assert_called_once_with("tenant123", exclude_vector_space=True) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") @patch("controllers.service_api.wraps.FeatureService.get_vector_space") def test_loads_vector_space_from_dedicated_quota_api( self, mock_get_vector_space, mock_get_features, mock_validate_token, app: Flask ): """Test vector-space resource checks avoid loading the full feature payload.""" # Arrange mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_vector_space = Mock() mock_vector_space.limit = 10 mock_vector_space.size = 5 mock_get_vector_space.return_value = mock_vector_space @cloud_edition_billing_resource_check("vector_space", "dataset") def add_segment(): return "segment_added" # Act with ( app.test_request_context("/", method="GET"), patch("controllers.service_api.wraps.dify_config.BILLING_ENABLED", True), ): result = add_segment() # Assert assert result == "segment_added" mock_get_vector_space.assert_called_once_with("tenant123") mock_get_features.assert_not_called() @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_loads_features_when_checking_non_vector_space_limit( self, mock_get_features, mock_validate_token, app: Flask ): """Test non-vector-space resource checks keep using the light feature payload.""" # Arrange mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() mock_features.billing.enabled = True mock_features.documents_upload_quota.limit = 10 mock_features.documents_upload_quota.size = 5 mock_get_features.return_value = mock_features @cloud_edition_billing_resource_check("documents", "dataset") def upload_document(): return "document_uploaded" # Act with app.test_request_context("/", method="GET"): result = upload_document() # Assert assert result == "document_uploaded" mock_get_features.assert_called_once_with("tenant123", exclude_vector_space=True) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_rejects_when_at_limit(self, mock_get_features, mock_validate_token, app: Flask): """Test that Forbidden is raised when at resource limit.""" # Arrange mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() mock_features.billing.enabled = True mock_features.members.limit = 10 mock_features.members.size = 10 mock_get_features.return_value = mock_features @cloud_edition_billing_resource_check("members", "app") def add_member(): return "member_added" # Act & Assert with app.test_request_context("/", method="GET"): with pytest.raises(Forbidden) as exc_info: add_member() assert "members has reached the limit" in str(exc_info.value) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_allows_when_billing_disabled(self, mock_get_features, mock_validate_token, app: Flask): """Test that request is allowed when billing is disabled.""" # Arrange mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() mock_features.billing.enabled = False mock_get_features.return_value = mock_features @cloud_edition_billing_resource_check("members", "app") def add_member(): return "member_added" # Act with app.test_request_context("/", method="GET"): result = add_member() # Assert assert result == "member_added" class TestCloudEditionBillingKnowledgeLimitCheck: """Test suite for cloud_edition_billing_knowledge_limit_check decorator""" @pytest.fixture def app(self): """Create Flask test application.""" app = Flask(__name__) app.config["TESTING"] = True return app @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_rejects_add_segment_in_sandbox(self, mock_get_features, mock_validate_token, app: Flask): """Test that add_segment is rejected in SANDBOX plan.""" # Arrange mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() mock_features.billing.enabled = True mock_features.billing.subscription.plan = CloudPlan.SANDBOX mock_get_features.return_value = mock_features @cloud_edition_billing_knowledge_limit_check("add_segment", "dataset") def add_segment(): return "segment_added" # Act & Assert with app.test_request_context("/", method="GET"): with pytest.raises(Forbidden) as exc_info: add_segment() assert "upgrade to a paid plan" in str(exc_info.value) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_allows_other_operations_in_sandbox(self, mock_get_features, mock_validate_token, app: Flask): """Test that non-add_segment operations are allowed in SANDBOX.""" # Arrange mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() mock_features.billing.enabled = True mock_features.billing.subscription.plan = CloudPlan.SANDBOX mock_get_features.return_value = mock_features @cloud_edition_billing_knowledge_limit_check("search", "dataset") def search(): return "search_results" # Act with app.test_request_context("/", method="GET"): result = search() # Assert assert result == "search_results" class TestCloudEditionBillingRateLimitCheck: """Test suite for cloud_edition_billing_rate_limit_check decorator""" @pytest.fixture def app(self): """Create Flask test application.""" app = Flask(__name__) app.config["TESTING"] = True return app @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_knowledge_rate_limit") def test_allows_within_rate_limit(self, mock_get_rate_limit, mock_validate_token, app: Flask): """Test that request is allowed when within rate limit.""" # Arrange mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_rate_limit = Mock() mock_rate_limit.enabled = True mock_rate_limit.limit = 100 mock_get_rate_limit.return_value = mock_rate_limit # Mock redis operations with patch("controllers.service_api.wraps.redis_client") as mock_redis: mock_redis.zcard.return_value = 50 # Under limit @cloud_edition_billing_rate_limit_check("knowledge", "dataset") def knowledge_request(): return "success" # Act with app.test_request_context("/", method="GET"): result = knowledge_request() # Assert assert result == "success" mock_redis.zadd.assert_called_once() mock_redis.zremrangebyscore.assert_called_once() @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_knowledge_rate_limit") @patch("controllers.service_api.wraps.db") @patch("controllers.service_api.wraps.sessionmaker") def test_rejects_over_rate_limit( self, mock_sessionmaker, mock_db, mock_get_rate_limit, mock_validate_token, app: Flask ): """Test that Forbidden is raised when over rate limit.""" # Arrange mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_rate_limit = Mock() mock_rate_limit.enabled = True mock_rate_limit.limit = 10 mock_rate_limit.subscription_plan = "pro" mock_get_rate_limit.return_value = mock_rate_limit rate_limit_log_session = MagicMock() session_factory = MagicMock() session_factory.begin.return_value.__enter__.return_value = rate_limit_log_session mock_sessionmaker.return_value = session_factory with patch("controllers.service_api.wraps.redis_client") as mock_redis: mock_redis.zcard.return_value = 15 # Over limit @cloud_edition_billing_rate_limit_check("knowledge", "dataset") def knowledge_request(): return "success" # Act & Assert with app.test_request_context("/", method="GET"): with pytest.raises(Forbidden) as exc_info: knowledge_request() assert "rate limit" in str(exc_info.value) mock_sessionmaker.assert_called_once_with(bind=mock_db.engine, expire_on_commit=False) rate_limit_log_session.add.assert_called_once() mock_db.session.commit.assert_not_called() class TestValidateDatasetToken: """Test suite for validate_dataset_token decorator""" @pytest.fixture def app(self): """Create Flask test application.""" app = Flask(__name__) app.config["TESTING"] = True return app @patch("controllers.service_api.wraps.user_logged_in") @patch("controllers.service_api.wraps.db") @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.current_app") def test_valid_dataset_token(self, mock_current_app, mock_validate_token, mock_db, mock_user_logged_in, app: Flask): """Test that valid dataset token allows access.""" # Arrange _configure_current_app_mock(mock_current_app) tenant_id = str(uuid.uuid4()) mock_api_token = Mock() mock_api_token.tenant_id = tenant_id mock_validate_token.return_value = mock_api_token mock_tenant = Mock() mock_tenant.id = tenant_id mock_tenant.status = TenantStatus.NORMAL mock_ta = Mock() mock_ta.account_id = str(uuid.uuid4()) mock_account = Mock() mock_account.id = mock_ta.account_id mock_account.current_tenant = mock_tenant # Mock the tenant account join query (execute(select(...)).one_or_none()) setup_mock_dataset_owner_execute_result(mock_db, mock_tenant, mock_ta) # Mock the account lookup via session.get() mock_db.session.get.return_value = mock_account @validate_dataset_token def protected_view(tenant_id): return {"success": True, "tenant_id": tenant_id} # Act with app.test_request_context("/", method="GET", headers={"Authorization": "Bearer test_token"}): result = protected_view() # Assert assert result["success"] is True assert result["tenant_id"] == tenant_id @patch("controllers.service_api.wraps.db") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_dataset_not_found_raises_not_found(self, mock_validate_token, mock_db, app: Flask): """Test that NotFound is raised when dataset doesn't exist.""" # Arrange mock_api_token = Mock() mock_api_token.tenant_id = str(uuid.uuid4()) mock_validate_token.return_value = mock_api_token mock_db.session.scalar.return_value = None @validate_dataset_token def protected_view(dataset_id=None, **kwargs): return {"success": True} # Act & Assert with app.test_request_context("/", method="GET"): with pytest.raises(NotFound) as exc_info: protected_view(dataset_id=str(uuid.uuid4())) assert "Dataset not found" in str(exc_info.value) class TestFetchUserArg: """Test suite for FetchUserArg model""" def test_fetch_user_arg_defaults(self): """Test FetchUserArg default values.""" # Arrange & Act arg = FetchUserArg(fetch_from=WhereisUserArg.JSON) # Assert assert arg.fetch_from == WhereisUserArg.JSON assert arg.required is False def test_fetch_user_arg_required(self): """Test FetchUserArg with required=True.""" # Arrange & Act arg = FetchUserArg(fetch_from=WhereisUserArg.QUERY, required=True) # Assert assert arg.fetch_from == WhereisUserArg.QUERY assert arg.required is True class TestDatasetApiResource: """Test suite for DatasetApiResource base class""" def test_method_decorators_has_validate_dataset_token(self): """Test that DatasetApiResource has validate_dataset_token in method_decorators.""" # Assert assert validate_dataset_token in DatasetApiResource.method_decorators def test_get_dataset_method_exists(self): """Test that get_dataset method exists on DatasetApiResource.""" # Assert assert hasattr(DatasetApiResource, "get_dataset")