mirror of
https://github.com/langgenius/dify.git
synced 2026-06-26 14:51:13 +08:00
be6612f454
Reintroduce the nullable api_tokens.dataset_id column (dropped in 2e9819ca5b28) so dataset API keys can opt into per-knowledge-base scoping: - NULL dataset_id keeps today's workspace-wide behavior, so every existing key and the existing /datasets/api-keys create route are unchanged. - validate_dataset_token rejects a bound key for any other dataset, and for endpoints that carry no dataset id (e.g. list-all), with 403. - CachedApiToken carries dataset_id with a None default so cache entries written before deploy keep deserializing. - The per-dataset console routes in apikey.py (previously dead code that 500ed on a missing ApiToken.dataset_id) now create bound keys; their list returns bound keys plus workspace keys so the dataset page shows the full access picture. - Frontend: the knowledge base API access popover gains an API keys entry; the secret key modal accepts datasetId, shows a scope column, and offers a workspace / this-knowledge-base scope choice on create. New strings are localized for all 23 locales.
736 lines
28 KiB
Python
736 lines
28 KiB
Python
"""
|
|
Unit tests for Service API wraps (authentication decorators)
|
|
"""
|
|
|
|
import uuid
|
|
from unittest.mock import 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")
|
|
def test_rejects_over_rate_limit(self, 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
|
|
|
|
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)
|
|
|
|
|
|
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_api_token.dataset_id = None
|
|
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_api_token.dataset_id = None
|
|
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)
|
|
|
|
def _arrange_tenant_owner(self, mock_db, tenant_id: str) -> None:
|
|
"""Stub the tenant-owner resolution shared by all success-path tests."""
|
|
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
|
|
|
|
setup_mock_dataset_owner_execute_result(mock_db, mock_tenant, mock_ta)
|
|
mock_db.session.get.return_value = mock_account
|
|
|
|
@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_workspace_scoped_token_can_access_any_dataset(
|
|
self, mock_current_app, mock_validate_token, mock_db, mock_user_logged_in, app: Flask
|
|
):
|
|
"""A token without dataset binding (NULL dataset_id) keeps tenant-wide access."""
|
|
# Arrange
|
|
_configure_current_app_mock(mock_current_app)
|
|
|
|
tenant_id = str(uuid.uuid4())
|
|
dataset_id = str(uuid.uuid4())
|
|
mock_api_token = Mock()
|
|
mock_api_token.tenant_id = tenant_id
|
|
mock_api_token.dataset_id = None
|
|
mock_validate_token.return_value = mock_api_token
|
|
|
|
mock_dataset = Mock()
|
|
mock_dataset.id = dataset_id
|
|
mock_dataset.enable_api = True
|
|
mock_db.session.scalar.return_value = mock_dataset
|
|
|
|
self._arrange_tenant_owner(mock_db, tenant_id)
|
|
|
|
@validate_dataset_token
|
|
def protected_view(tenant_id, dataset_id=None):
|
|
return {"success": True}
|
|
|
|
# Act
|
|
with app.test_request_context("/", method="GET", headers={"Authorization": "Bearer test_token"}):
|
|
result = protected_view(dataset_id=dataset_id)
|
|
|
|
# Assert
|
|
assert result["success"] is True
|
|
|
|
@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_dataset_bound_token_allows_its_own_dataset(
|
|
self, mock_current_app, mock_validate_token, mock_db, mock_user_logged_in, app: Flask
|
|
):
|
|
"""A dataset-bound token can access the dataset it is bound to."""
|
|
# Arrange
|
|
_configure_current_app_mock(mock_current_app)
|
|
|
|
tenant_id = str(uuid.uuid4())
|
|
dataset_id = str(uuid.uuid4())
|
|
mock_api_token = Mock()
|
|
mock_api_token.tenant_id = tenant_id
|
|
mock_api_token.dataset_id = dataset_id
|
|
mock_validate_token.return_value = mock_api_token
|
|
|
|
mock_dataset = Mock()
|
|
mock_dataset.id = dataset_id
|
|
mock_dataset.enable_api = True
|
|
mock_db.session.scalar.return_value = mock_dataset
|
|
|
|
self._arrange_tenant_owner(mock_db, tenant_id)
|
|
|
|
@validate_dataset_token
|
|
def protected_view(tenant_id, dataset_id=None):
|
|
return {"success": True}
|
|
|
|
# Act
|
|
with app.test_request_context("/", method="GET", headers={"Authorization": "Bearer test_token"}):
|
|
result = protected_view(dataset_id=dataset_id)
|
|
|
|
# Assert
|
|
assert result["success"] is True
|
|
|
|
@patch("controllers.service_api.wraps.db")
|
|
@patch("controllers.service_api.wraps.validate_and_get_api_token")
|
|
def test_dataset_bound_token_rejects_other_dataset(self, mock_validate_token, mock_db, app: Flask):
|
|
"""A dataset-bound token gets Forbidden for any other dataset in the tenant."""
|
|
# Arrange
|
|
mock_api_token = Mock()
|
|
mock_api_token.tenant_id = str(uuid.uuid4())
|
|
mock_api_token.dataset_id = str(uuid.uuid4())
|
|
mock_validate_token.return_value = mock_api_token
|
|
|
|
@validate_dataset_token
|
|
def protected_view(tenant_id, dataset_id=None):
|
|
return {"success": True}
|
|
|
|
# Act & Assert
|
|
with app.test_request_context("/", method="GET"):
|
|
with pytest.raises(Forbidden) as exc_info:
|
|
protected_view(dataset_id=str(uuid.uuid4()))
|
|
assert "not authorized" in str(exc_info.value)
|
|
|
|
@patch("controllers.service_api.wraps.db")
|
|
@patch("controllers.service_api.wraps.validate_and_get_api_token")
|
|
def test_dataset_bound_token_rejects_request_without_dataset_id(self, mock_validate_token, mock_db, app: Flask):
|
|
"""A dataset-bound token cannot call endpoints that carry no dataset id (e.g. list-all)."""
|
|
# Arrange
|
|
mock_api_token = Mock()
|
|
mock_api_token.tenant_id = str(uuid.uuid4())
|
|
mock_api_token.dataset_id = str(uuid.uuid4())
|
|
mock_validate_token.return_value = mock_api_token
|
|
|
|
@validate_dataset_token
|
|
def protected_view(tenant_id):
|
|
return {"success": True}
|
|
|
|
# Act & Assert
|
|
with app.test_request_context("/", method="GET"):
|
|
with pytest.raises(Forbidden) as exc_info:
|
|
protected_view()
|
|
assert "not authorized" 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")
|