test: 补充知识处理、检索与 Agent 回归用例
覆盖切块、向量检索、工具调用和历史预算;模型相关用例使用桩。
This commit is contained in:
1 parent
2dd062787d
commit
d9fe727c3f
123 files changed
+29005
No files matched your search
@@ -0,0 +1,16 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
|
||||
agno reports each agent run to api.agno.com from inside the run unless told
|
||||
not to; a slow report used to eat the run's timeout and cost the visitor
|
||||
their answer.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from agno.agent import Agent
|
||||
from agno.run.response import RunResponse
|
||||
|
||||
import app # noqa: F401 — app/__init__.py, which every process imports first, applies the switches
|
||||
from app.core import agno_telemetry
|
||||
from app.core.config import Settings, settings
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def configured(monkeypatch):
|
||||
"""Set the switches as an operator would, and restore the process env after."""
|
||||
for name in (*agno_telemetry.AGNO_SWITCHES, agno_telemetry.AGNO_API_SWITCH):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
|
||||
def configure(**switches):
|
||||
for name in agno_telemetry.AGNO_SWITCHES:
|
||||
monkeypatch.setattr(settings, name, switches.get(name, False))
|
||||
agno_telemetry.apply()
|
||||
|
||||
return configure
|
||||
|
||||
|
||||
def test_off_by_default():
|
||||
assert Settings.model_fields["AGNO_TELEMETRY"].default is False
|
||||
assert Settings.model_fields["AGNO_MONITOR"].default is False
|
||||
|
||||
|
||||
def test_apply_turns_agnos_reporting_off(configured):
|
||||
configured()
|
||||
|
||||
assert os.environ["AGNO_TELEMETRY"] == "false"
|
||||
assert os.environ["AGNO_MONITOR"] == "false"
|
||||
assert os.environ["AGNO_API_ENABLED"] == "false"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("switch", agno_telemetry.AGNO_SWITCHES)
|
||||
def test_an_operator_can_opt_in(configured, switch):
|
||||
configured(**{switch: True})
|
||||
|
||||
assert os.environ[switch] == "true"
|
||||
# Their opt-in must still reach agno's API.
|
||||
assert agno_telemetry.AGNO_API_SWITCH not in os.environ
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value, expected", [("true", True), ("TRUE", True), ("false", False), ("0", False)])
|
||||
def test_the_setting_comes_from_the_environment_or_env_file(monkeypatch, value, expected):
|
||||
# backend/.env reaches the environment before Settings reads it (config.py's load_dotenv).
|
||||
monkeypatch.setenv("AGNO_TELEMETRY", value)
|
||||
|
||||
assert Settings().AGNO_TELEMETRY is expected
|
||||
|
||||
|
||||
def test_the_app_applied_it_on_import():
|
||||
for name in agno_telemetry.AGNO_SWITCHES:
|
||||
assert os.environ[name] == ("true" if getattr(settings, name) else "false")
|
||||
|
||||
|
||||
def _finished_agent() -> Agent:
|
||||
"""An agno Agent at the point it reports a run: its response is in hand."""
|
||||
agent = Agent()
|
||||
agent.run_response = RunResponse(content="done")
|
||||
return agent
|
||||
|
||||
|
||||
async def _reports_sent(agent: Agent) -> int:
|
||||
with patch("agno.api.agent.acreate_agent_run", new_callable=AsyncMock) as report, \
|
||||
patch("agno.api.agent.create_agent_run") as sync_report:
|
||||
await agent._alog_agent_run(session_id="session-1")
|
||||
agent._log_agent_run(session_id="session-1")
|
||||
return report.await_count + sync_report.call_count
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_agent_built_with_agnos_defaults_reports_nothing(configured):
|
||||
"""Guards the env names against an agno upgrade: Agent() still defaults to
|
||||
telemetry=True, and only these variables switch it off at run time."""
|
||||
configured()
|
||||
agent = _finished_agent()
|
||||
assert agent.telemetry is True # agno's own default, overridden per run
|
||||
|
||||
assert await _reports_sent(agent) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_opted_in_agent_does_report(configured):
|
||||
"""The other half of the guard: with telemetry on, the report does go out -
|
||||
so a zero above means switched off, not a test that can't see it."""
|
||||
configured(AGNO_TELEMETRY=True)
|
||||
|
||||
assert await _reports_sent(_finished_agent()) == 2
|
||||
@@ -0,0 +1,880 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
from fastapi import HTTPException, Request, Depends
|
||||
from sqlalchemy.orm import Session
|
||||
from uuid import uuid4
|
||||
|
||||
from app.core.auth import (
|
||||
check_permissions,
|
||||
get_current_user,
|
||||
get_current_organization,
|
||||
require_permissions,
|
||||
require_permission,
|
||||
require_subscription_management,
|
||||
require_unified_permissions,
|
||||
require_any_unified_permission
|
||||
)
|
||||
from app.models.user import User
|
||||
from app.models.organization import Organization
|
||||
from app.models.permission import Permission
|
||||
from app.models.role import Role
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_db():
|
||||
"""Create a mock database session"""
|
||||
db = MagicMock(spec=Session)
|
||||
return db
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_user():
|
||||
"""Create a mock user with role and permissions"""
|
||||
user_id = uuid4()
|
||||
org_id = uuid4()
|
||||
|
||||
# Create permissions
|
||||
view_chats_permission = Permission(id=1, name="view_chats")
|
||||
manage_agents_permission = Permission(id=2, name="manage_agents")
|
||||
|
||||
# Create role with permissions
|
||||
role = Role(id=1, name="Agent", permissions=[view_chats_permission, manage_agents_permission])
|
||||
|
||||
# Create user with role
|
||||
user = User(
|
||||
id=user_id,
|
||||
email="test@example.com",
|
||||
hashed_password="hashed_password",
|
||||
is_active=True,
|
||||
organization_id=org_id,
|
||||
full_name="Test User",
|
||||
role=role
|
||||
)
|
||||
|
||||
return user
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_super_admin():
|
||||
"""Create a mock user with super_admin permission"""
|
||||
user_id = uuid4()
|
||||
org_id = uuid4()
|
||||
|
||||
# Create super_admin permission
|
||||
super_admin_permission = Permission(id=3, name="super_admin")
|
||||
|
||||
# Create role with super_admin permission
|
||||
role = Role(id=2, name="Super Admin", permissions=[super_admin_permission])
|
||||
|
||||
# Create user with role
|
||||
user = User(
|
||||
id=user_id,
|
||||
email="admin@example.com",
|
||||
hashed_password="hashed_password",
|
||||
is_active=True,
|
||||
organization_id=org_id,
|
||||
full_name="Admin User",
|
||||
role=role
|
||||
)
|
||||
|
||||
return user
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_organization():
|
||||
"""Create a mock organization"""
|
||||
org_id = uuid4()
|
||||
return Organization(
|
||||
id=org_id,
|
||||
name="Test Organization",
|
||||
domain="test.com",
|
||||
is_active=True
|
||||
)
|
||||
|
||||
|
||||
def test_check_permissions_with_required_permissions(mock_user):
|
||||
"""Test check_permissions when user has the required permissions"""
|
||||
result = check_permissions(mock_user, ["view_chats"])
|
||||
assert result is True
|
||||
|
||||
|
||||
def test_check_permissions_without_required_permissions(mock_user):
|
||||
"""Test check_permissions when user doesn't have the required permissions"""
|
||||
result = check_permissions(mock_user, ["delete_organization"])
|
||||
assert result is False
|
||||
|
||||
|
||||
def test_check_permissions_with_super_admin(mock_super_admin):
|
||||
"""Test check_permissions with super_admin permission"""
|
||||
result = check_permissions(mock_super_admin, ["delete_organization"])
|
||||
assert result is True
|
||||
|
||||
|
||||
def test_check_permissions_without_role():
|
||||
"""Test check_permissions when user has no role"""
|
||||
user = User(
|
||||
id=uuid4(),
|
||||
email="noauth@example.com",
|
||||
hashed_password="hashed_password",
|
||||
is_active=True
|
||||
)
|
||||
result = check_permissions(user, ["view_chats"])
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_user_with_cookie(mock_db, mock_user):
|
||||
"""Test get_current_user with access token in cookie"""
|
||||
# Mock request with cookie
|
||||
request = MagicMock(spec=Request)
|
||||
access_token = "valid_access_token"
|
||||
|
||||
# Mock live_payload
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
|
||||
# Mock db query
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
result = await get_current_user(request, access_token, mock_db)
|
||||
assert result == mock_user
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_user_with_auth_header(mock_db, mock_user):
|
||||
"""Test get_current_user with access token in Authorization header"""
|
||||
# Mock request with Authorization header
|
||||
request = MagicMock(spec=Request)
|
||||
request.headers.get.return_value = "Bearer valid_access_token"
|
||||
|
||||
# Mock live_payload
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
|
||||
# Mock db query
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
result = await get_current_user(request, None, mock_db)
|
||||
assert result == mock_user
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_user_no_token(mock_db):
|
||||
"""Test get_current_user with no access token"""
|
||||
# Mock request with no token
|
||||
request = MagicMock(spec=Request)
|
||||
request.headers.get.return_value = None
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_current_user(request, None, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail == "Not authenticated"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_user_invalid_token(mock_db):
|
||||
"""Test get_current_user with invalid token"""
|
||||
# Mock request
|
||||
request = MagicMock(spec=Request)
|
||||
access_token = "invalid_token"
|
||||
|
||||
# Mock live_payload
|
||||
with patch('app.core.auth.live_payload', return_value=None):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_current_user(request, access_token, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail == "Invalid authentication token"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_user_invalid_payload(mock_db):
|
||||
"""Test get_current_user with invalid payload"""
|
||||
# Mock request
|
||||
request = MagicMock(spec=Request)
|
||||
access_token = "valid_token_invalid_payload"
|
||||
|
||||
# Mock live_payload
|
||||
payload = {"not_sub": "invalid"}
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_current_user(request, access_token, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail == "Invalid token payload"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_user_user_not_found(mock_db):
|
||||
"""Test get_current_user when user not found"""
|
||||
# Mock request
|
||||
request = MagicMock(spec=Request)
|
||||
access_token = "valid_token"
|
||||
|
||||
# Mock live_payload
|
||||
payload = {"sub": str(uuid4())}
|
||||
|
||||
# Mock db query - user not found
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = None
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_current_user(request, access_token, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail == "User not found"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_user_inactive_user(mock_db, mock_user):
|
||||
"""Test get_current_user with inactive user"""
|
||||
# Mock request
|
||||
request = MagicMock(spec=Request)
|
||||
access_token = "valid_token"
|
||||
|
||||
# Set user as inactive
|
||||
mock_user.is_active = False
|
||||
|
||||
# Mock live_payload
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
|
||||
# Mock db query
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_current_user(request, access_token, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail == "User is inactive"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_user_exception(mock_db):
|
||||
"""Test get_current_user with exception"""
|
||||
# Mock request
|
||||
request = MagicMock(spec=Request)
|
||||
access_token = "valid_token"
|
||||
|
||||
# Mock live_payload to raise exception
|
||||
with patch('app.core.auth.live_payload', side_effect=Exception("Test exception")):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_current_user(request, access_token, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail == "Could not validate credentials"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_organization_success(mock_db, mock_user, mock_organization):
|
||||
"""Test get_current_organization success"""
|
||||
# Mock db query
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_organization
|
||||
|
||||
result = await get_current_organization(mock_user, mock_db)
|
||||
assert result == mock_organization
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_organization_no_org_id(mock_db):
|
||||
"""Test get_current_organization with no organization_id"""
|
||||
# Create user with no organization_id
|
||||
user = User(
|
||||
id=uuid4(),
|
||||
email="noorg@example.com",
|
||||
hashed_password="hashed_password",
|
||||
is_active=True,
|
||||
organization_id=None
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_current_organization(user, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == "User is not associated with any organization"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_organization_not_found(mock_db, mock_user):
|
||||
"""Test get_current_organization when organization not found"""
|
||||
# Mock db query - organization not found
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = None
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_current_organization(mock_user, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert exc_info.value.detail == "Organization not found or inactive"
|
||||
|
||||
|
||||
# Test the require_permissions function indirectly through check_permissions
|
||||
def test_require_permissions_check_with_permissions(mock_user):
|
||||
"""Test the check used in require_permissions when user has the required permissions"""
|
||||
result = check_permissions(mock_user, ["view_chats"])
|
||||
assert result is True
|
||||
|
||||
|
||||
def test_require_permissions_check_without_permissions(mock_user):
|
||||
"""Test the check used in require_permissions when user doesn't have the required permissions"""
|
||||
result = check_permissions(mock_user, ["delete_organization"])
|
||||
assert result is False
|
||||
|
||||
|
||||
# Test the require_permission function indirectly through role.has_permission
|
||||
def test_require_permission_check_with_permission(mock_user):
|
||||
"""Test the check used in require_permission when user has the required permission"""
|
||||
# Mock the has_permission method
|
||||
mock_user.role.has_permission = MagicMock(return_value=True)
|
||||
|
||||
# Check if the role has the permission
|
||||
result = mock_user.role.has_permission("view_chats")
|
||||
assert result is True
|
||||
|
||||
# Verify the has_permission method was called with the correct permission
|
||||
mock_user.role.has_permission.assert_called_once_with("view_chats")
|
||||
|
||||
|
||||
def test_require_permission_check_without_permission(mock_user):
|
||||
"""Test the check used in require_permission when user doesn't have the required permission"""
|
||||
# Mock the has_permission method
|
||||
mock_user.role.has_permission = MagicMock(return_value=False)
|
||||
|
||||
# Check if the role has the permission
|
||||
result = mock_user.role.has_permission("delete_organization")
|
||||
assert result is False
|
||||
|
||||
# Verify the has_permission method was called with the correct permission
|
||||
mock_user.role.has_permission.assert_called_once_with("delete_organization")
|
||||
|
||||
|
||||
# Test the require_subscription_management function indirectly through role.can_manage_subscription
|
||||
def test_require_subscription_management_check_with_permission(mock_user):
|
||||
"""Test the check used in require_subscription_management when user has the permission"""
|
||||
# Mock the can_manage_subscription method
|
||||
mock_user.role.can_manage_subscription = MagicMock(return_value=True)
|
||||
|
||||
# Check if the role can manage subscriptions
|
||||
result = mock_user.role.can_manage_subscription()
|
||||
assert result is True
|
||||
|
||||
# Verify the can_manage_subscription method was called
|
||||
mock_user.role.can_manage_subscription.assert_called_once()
|
||||
|
||||
|
||||
def test_require_subscription_management_check_without_permission(mock_user):
|
||||
"""Test the check used in require_subscription_management when user doesn't have the permission"""
|
||||
# Mock the can_manage_subscription method
|
||||
mock_user.role.can_manage_subscription = MagicMock(return_value=False)
|
||||
|
||||
# Check if the role can manage subscriptions
|
||||
result = mock_user.role.can_manage_subscription()
|
||||
assert result is False
|
||||
|
||||
# Verify the can_manage_subscription method was called
|
||||
mock_user.role.can_manage_subscription.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_require_permissions_inner_function(mock_user):
|
||||
"""Test the inner function of require_permissions directly"""
|
||||
# Create the inner function
|
||||
permission_checker = require_permissions("view_chats")
|
||||
|
||||
# Mock Depends to return our user
|
||||
with patch('fastapi.Depends', return_value=mock_user):
|
||||
# Call the inner function
|
||||
result = await permission_checker(mock_user)
|
||||
assert result == mock_user
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_require_permissions_inner_function_forbidden(mock_user):
|
||||
"""Test the inner function of require_permissions with insufficient permissions"""
|
||||
# Create the inner function
|
||||
permission_checker = require_permissions("delete_organization")
|
||||
|
||||
# Call the inner function and expect an exception
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await permission_checker(mock_user)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == "Not enough permissions"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_require_permission_inner_function(mock_user):
|
||||
"""Test the inner function of require_permission directly"""
|
||||
# Mock the role's has_permission method
|
||||
mock_user.role.has_permission = MagicMock(return_value=True)
|
||||
|
||||
# Create the inner function
|
||||
permission_dependency = require_permission("view_chats")
|
||||
|
||||
# Call the inner function
|
||||
result = await permission_dependency(mock_user)
|
||||
assert result == mock_user
|
||||
mock_user.role.has_permission.assert_called_once_with("view_chats")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_require_permission_inner_function_forbidden(mock_user):
|
||||
"""Test the inner function of require_permission with insufficient permissions"""
|
||||
# Mock the role's has_permission method
|
||||
mock_user.role.has_permission = MagicMock(return_value=False)
|
||||
|
||||
# Create the inner function
|
||||
permission_dependency = require_permission("delete_organization")
|
||||
|
||||
# Call the inner function and expect an exception
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await permission_dependency(mock_user)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == "User does not have required permission: delete_organization"
|
||||
mock_user.role.has_permission.assert_called_once_with("delete_organization")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_require_subscription_management_with_permission(mock_user):
|
||||
"""Test require_subscription_management with sufficient permissions"""
|
||||
# Mock the role's can_manage_subscription method
|
||||
mock_user.role.can_manage_subscription = MagicMock(return_value=True)
|
||||
|
||||
# Call the function
|
||||
result = require_subscription_management(mock_user)
|
||||
assert result == mock_user
|
||||
mock_user.role.can_manage_subscription.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_require_subscription_management_without_permission(mock_user):
|
||||
"""Test require_subscription_management with insufficient permissions"""
|
||||
# Mock the role's can_manage_subscription method
|
||||
mock_user.role.can_manage_subscription = MagicMock(return_value=False)
|
||||
|
||||
# Call the function and expect an exception
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
require_subscription_management(mock_user)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == "User does not have permission to manage subscriptions"
|
||||
mock_user.role.can_manage_subscription.assert_called_once()
|
||||
|
||||
|
||||
# ==================== Tests for get_auth_info_from_request ====================
|
||||
|
||||
from app.core.auth import get_auth_info_from_request
|
||||
|
||||
def test_get_auth_info_from_request_shopify_url():
|
||||
"""Test get_auth_info_from_request with Shopify URL"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/shopify/products"
|
||||
|
||||
result = get_auth_info_from_request(request)
|
||||
assert result["is_shopify"] is True
|
||||
|
||||
|
||||
def test_get_auth_info_from_request_non_shopify_url():
|
||||
"""Test get_auth_info_from_request with non-Shopify URL"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/users"
|
||||
|
||||
result = get_auth_info_from_request(request)
|
||||
assert result["is_shopify"] is False
|
||||
|
||||
|
||||
# ==================== Tests for get_unified_auth ====================
|
||||
|
||||
from app.core.auth import get_unified_auth
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_auth_jwt_success(mock_db, mock_user):
|
||||
"""Test get_unified_auth with valid JWT token"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/agents"
|
||||
request.query_params = {}
|
||||
request.cookies.get.return_value = "valid_token"
|
||||
request.headers.get.return_value = None
|
||||
|
||||
# Add manage_agents permission
|
||||
manage_agents_perm = Permission(id=2, name="manage_agents")
|
||||
mock_user.role.permissions = [manage_agents_perm]
|
||||
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
result = await get_unified_auth(request, mock_db)
|
||||
|
||||
assert result["auth_type"] == "jwt"
|
||||
assert result["organization_id"] == mock_user.organization_id
|
||||
assert result["current_user"] == mock_user
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_auth_jwt_no_token(mock_db):
|
||||
"""Test get_unified_auth without token"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/agents"
|
||||
request.query_params = {}
|
||||
request.cookies.get.return_value = None
|
||||
request.headers.get.return_value = None
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_unified_auth(request, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_auth_jwt_invalid_token(mock_db):
|
||||
"""Test get_unified_auth with invalid JWT token"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/agents"
|
||||
request.query_params = {}
|
||||
request.cookies.get.return_value = "invalid_token"
|
||||
request.headers.get.return_value = None
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=None):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_unified_auth(request, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail == "Invalid authentication token"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_auth_jwt_user_not_found(mock_db):
|
||||
"""Test get_unified_auth when user not found"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/agents"
|
||||
request.query_params = {}
|
||||
request.cookies.get.return_value = "valid_token"
|
||||
request.headers.get.return_value = None
|
||||
|
||||
payload = {"sub": str(uuid4())}
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = None
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_unified_auth(request, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail == "User not found"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_auth_jwt_inactive_user(mock_db, mock_user):
|
||||
"""Test get_unified_auth with inactive user"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/agents"
|
||||
request.query_params = {}
|
||||
request.cookies.get.return_value = "valid_token"
|
||||
request.headers.get.return_value = None
|
||||
|
||||
mock_user.is_active = False
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_unified_auth(request, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.detail == "User is inactive"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_auth_authenticates_without_permissions(mock_db, mock_user):
|
||||
"""get_unified_auth authenticates; it does not authorize.
|
||||
|
||||
It used to hardcode a manage_agents check, so every route sharing the helper
|
||||
inherited it — including reading your own org's subscription, which locked
|
||||
every non-admin out of the entire app.
|
||||
"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/agents"
|
||||
request.query_params = {}
|
||||
request.cookies.get.return_value = "valid_token"
|
||||
request.headers.get.return_value = None
|
||||
|
||||
mock_user.role.permissions = [Permission(id=1, name="view_only")]
|
||||
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
result = await get_unified_auth(request, mock_db)
|
||||
|
||||
assert result["auth_type"] == "jwt"
|
||||
assert result["current_user"] is mock_user
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_require_unified_permissions_rejects_missing_permission(mock_user):
|
||||
"""The wrapper is what enforces the route's requirement."""
|
||||
mock_user.role.permissions = [Permission(id=1, name="view_only")]
|
||||
|
||||
checker = require_unified_permissions("manage_agents")
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await checker(auth_info={"auth_type": "jwt", "current_user": mock_user})
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == "Not enough permissions"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_require_unified_permissions_allows_holder(mock_user):
|
||||
"""A user holding the permission passes through unchanged."""
|
||||
mock_user.role.permissions = [Permission(id=1, name="manage_agents")]
|
||||
auth_info = {"auth_type": "jwt", "current_user": mock_user}
|
||||
|
||||
checker = require_unified_permissions("manage_agents")
|
||||
|
||||
assert await checker(auth_info=auth_info) is auth_info
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_require_any_unified_permission_is_or_semantics(mock_user):
|
||||
"""Any one of the listed permissions is enough."""
|
||||
mock_user.role.permissions = [Permission(id=1, name="view_agents")]
|
||||
|
||||
checker = require_any_unified_permission("view_agents", "manage_agents")
|
||||
result = await checker(auth_info={"auth_type": "jwt", "current_user": mock_user})
|
||||
|
||||
assert result["auth_type"] == "jwt"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unified_permissions_pass_through_shop_session():
|
||||
"""A Shopify shop session authenticates a store, not a person.
|
||||
|
||||
It carries no current_user and no role, so there is nothing to check — it
|
||||
must pass through exactly as it did before authorization moved out.
|
||||
"""
|
||||
auth_info = {"auth_type": "shopify", "organization_id": uuid4()}
|
||||
|
||||
assert await require_unified_permissions("manage_agents")(auth_info=auth_info) is auth_info
|
||||
assert await require_any_unified_permission("manage_agents")(auth_info=auth_info) is auth_info
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unified_permissions_tolerate_absent_current_user_key():
|
||||
"""Callers that omit the key entirely must not blow up with a KeyError."""
|
||||
auth_info = {"auth_type": "jwt", "organization_id": uuid4()}
|
||||
|
||||
assert await require_unified_permissions("manage_agents")(auth_info=auth_info) is auth_info
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_auth_shopify_context(mock_db):
|
||||
"""Test get_unified_auth with Shopify context"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/shopify/agents"
|
||||
request.query_params = {}
|
||||
|
||||
mock_auth_result = {
|
||||
"auth_type": "shopify",
|
||||
"organization_id": uuid4(),
|
||||
"user_id": uuid4(),
|
||||
"current_user": None
|
||||
}
|
||||
|
||||
with patch('app.services.shopify_session.require_shopify_or_jwt_auth', new_callable=AsyncMock) as mock_shopify_auth:
|
||||
mock_shopify_auth.return_value = mock_auth_result
|
||||
result = await get_unified_auth(request, mock_db)
|
||||
|
||||
assert result["auth_type"] == "shopify"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_auth_with_auth_header(mock_db, mock_user):
|
||||
"""Test get_unified_auth with Authorization header"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/agents"
|
||||
request.query_params = {}
|
||||
request.cookies.get.return_value = None
|
||||
request.headers.get.return_value = "Bearer valid_token"
|
||||
|
||||
manage_agents_perm = Permission(id=2, name="manage_agents")
|
||||
mock_user.role.permissions = [manage_agents_perm]
|
||||
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
result = await get_unified_auth(request, mock_db)
|
||||
|
||||
assert result["auth_type"] == "jwt"
|
||||
assert result["current_user"] == mock_user
|
||||
|
||||
|
||||
# ==================== Tests for get_unified_chat_auth ====================
|
||||
|
||||
from app.core.auth import get_unified_chat_auth
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_chat_auth_jwt_success(mock_db, mock_user):
|
||||
"""Test get_unified_chat_auth with valid JWT token"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/chats"
|
||||
request.cookies.get.return_value = "valid_token"
|
||||
request.headers.get.return_value = None
|
||||
|
||||
# Add chat permissions
|
||||
view_all_chats = Permission(id=10, name="view_all_chats")
|
||||
mock_user.role.permissions = [view_all_chats]
|
||||
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
result = await get_unified_chat_auth(request, mock_db)
|
||||
|
||||
assert result["auth_type"] == "jwt"
|
||||
assert result["can_view_all"] is True
|
||||
assert result["can_view_assigned"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_chat_auth_view_assigned(mock_db, mock_user):
|
||||
"""Test get_unified_chat_auth with view_assigned_chats permission"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/chats"
|
||||
request.cookies.get.return_value = "valid_token"
|
||||
request.headers.get.return_value = None
|
||||
|
||||
# Add assigned chats permission only
|
||||
view_assigned_chats = Permission(id=11, name="view_assigned_chats")
|
||||
mock_user.role.permissions = [view_assigned_chats]
|
||||
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
result = await get_unified_chat_auth(request, mock_db)
|
||||
|
||||
assert result["auth_type"] == "jwt"
|
||||
assert result["can_view_all"] is False
|
||||
assert result["can_view_assigned"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_chat_auth_no_permissions(mock_db, mock_user):
|
||||
"""Test get_unified_chat_auth without chat permissions"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/chats"
|
||||
request.cookies.get.return_value = "valid_token"
|
||||
request.headers.get.return_value = None
|
||||
|
||||
# No chat permissions
|
||||
other_perm = Permission(id=1, name="other_permission")
|
||||
mock_user.role.permissions = [other_perm]
|
||||
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_unified_chat_auth(request, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == "Not enough permissions"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_chat_auth_shopify_context(mock_db):
|
||||
"""Test get_unified_chat_auth with Shopify context"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/shopify/chats"
|
||||
|
||||
mock_auth_result = {
|
||||
"auth_type": "shopify",
|
||||
"organization_id": uuid4(),
|
||||
"user_id": uuid4(),
|
||||
"current_user": None
|
||||
}
|
||||
|
||||
with patch('app.services.shopify_session.require_shopify_or_jwt_auth', new_callable=AsyncMock) as mock_shopify_auth:
|
||||
mock_shopify_auth.return_value = mock_auth_result
|
||||
result = await get_unified_chat_auth(request, mock_db)
|
||||
|
||||
assert result["auth_type"] == "shopify"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_chat_auth_no_token(mock_db):
|
||||
"""Test get_unified_chat_auth without token"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/chats"
|
||||
request.cookies.get.return_value = None
|
||||
request.headers.get.return_value = None
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_unified_chat_auth(request, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_chat_auth_exception(mock_db):
|
||||
"""Test get_unified_chat_auth with exception"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/chats"
|
||||
request.cookies.get.return_value = "valid_token"
|
||||
request.headers.get.return_value = None
|
||||
|
||||
with patch('app.core.auth.live_payload', side_effect=Exception("Test error")):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_unified_chat_auth(request, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_auth_exception(mock_db):
|
||||
"""Test get_unified_auth with unexpected exception"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.url = "https://example.com/api/agents"
|
||||
request.query_params = {}
|
||||
request.cookies.get.return_value = "valid_token"
|
||||
request.headers.get.return_value = None
|
||||
|
||||
with patch('app.core.auth.live_payload', side_effect=Exception("Unexpected error")):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_unified_auth(request, mock_db)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_user_from_cookie_directly(mock_db, mock_user):
|
||||
"""Test get_current_user when token is read directly from cookies"""
|
||||
request = MagicMock(spec=Request)
|
||||
request.cookies = {"access_token": "valid_cookie_token"}
|
||||
request.headers.get.return_value = None
|
||||
|
||||
payload = {"sub": str(mock_user.id)}
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_user
|
||||
|
||||
with patch('app.core.auth.live_payload', return_value=payload):
|
||||
result = await get_current_user(request, None, mock_db)
|
||||
assert result == mock_user
|
||||
@@ -0,0 +1,407 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
from app.core.auth_utils import (
|
||||
refresh_access_token,
|
||||
authenticate_socket,
|
||||
authenticate_socket_conversation_token,
|
||||
widget_socket_identity,
|
||||
)
|
||||
from app.models.user import User
|
||||
from app.models.widget import Widget
|
||||
from uuid import uuid4
|
||||
import http.cookies
|
||||
from datetime import datetime, timezone
|
||||
from app.core.security import verify_token, verify_conversation_token, create_access_token, create_refresh_token, create_conversation_token
|
||||
from app.models.auth_session import AuthSession
|
||||
from app.services.auth_sessions import issue_tokens, revoke_session
|
||||
|
||||
@pytest.fixture
|
||||
def mock_db():
|
||||
"""Create a mock database session"""
|
||||
db = MagicMock()
|
||||
return db
|
||||
|
||||
@pytest.fixture
|
||||
def mock_user():
|
||||
"""Create a mock user"""
|
||||
user_id = uuid4()
|
||||
org_id = uuid4()
|
||||
return User(
|
||||
id=user_id,
|
||||
email="test@example.com",
|
||||
hashed_password="hashed_password",
|
||||
is_active=True,
|
||||
organization_id=org_id,
|
||||
full_name="Test User"
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_widget():
|
||||
"""Create a mock widget"""
|
||||
widget_id = uuid4()
|
||||
org_id = uuid4()
|
||||
return Widget(
|
||||
id=str(widget_id),
|
||||
organization_id=org_id,
|
||||
name="Test Widget"
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def live_user(db, test_organization):
|
||||
"""A real, active user — session checks run against the database."""
|
||||
user = User(
|
||||
id=uuid4(),
|
||||
email="socket@example.com",
|
||||
hashed_password="hashed_password",
|
||||
is_active=True,
|
||||
organization_id=test_organization.id,
|
||||
full_name="Socket User",
|
||||
)
|
||||
db.add(user)
|
||||
db.commit()
|
||||
db.refresh(user)
|
||||
return user
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def socket_db(db):
|
||||
"""Point authenticate_socket's SessionLocal() at the test database."""
|
||||
with patch('app.core.auth_utils.SessionLocal') as mock_session_local:
|
||||
mock_session_local.return_value.__enter__.return_value = db
|
||||
yield db
|
||||
|
||||
|
||||
def _socket_environ(**cookies):
|
||||
# A same-origin handshake: only our own pages may sign in with the cookies
|
||||
jar = http.cookies.SimpleCookie()
|
||||
for name, value in cookies.items():
|
||||
jar[name] = value
|
||||
return {'HTTP_COOKIE': jar.output(header='', sep=';'), 'HTTP_SEC_FETCH_SITE': 'same-origin'}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_access_token_success(db, live_user):
|
||||
"""A live refresh token re-mints an access token for the same session"""
|
||||
_, refresh_token = issue_tokens(db, live_user.id, live_user.organization_id)
|
||||
db.commit()
|
||||
|
||||
result = await refresh_access_token(refresh_token, db)
|
||||
|
||||
payload = verify_token(result)
|
||||
assert payload["type"] == "access"
|
||||
assert payload["sub"] == str(live_user.id)
|
||||
assert payload["sid"] == verify_token(refresh_token)["sid"]
|
||||
assert db.query(AuthSession).count() == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_access_token_invalid_token(db):
|
||||
"""Test access token refresh with invalid token"""
|
||||
assert await refresh_access_token("invalid_token", db) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_access_token_inactive_user(db, live_user):
|
||||
"""Test access token refresh with inactive user"""
|
||||
_, refresh_token = issue_tokens(db, live_user.id, live_user.organization_id)
|
||||
live_user.is_active = False
|
||||
db.commit()
|
||||
|
||||
assert await refresh_access_token(refresh_token, db) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_access_token_rejects_revoked_session(db, live_user):
|
||||
"""A refresh token from a signed-out session mints nothing"""
|
||||
_, refresh_token = issue_tokens(db, live_user.id, live_user.organization_id)
|
||||
revoke_session(db, verify_token(refresh_token)["sid"])
|
||||
db.commit()
|
||||
|
||||
assert await refresh_access_token(refresh_token, db) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_access_token_rejects_an_access_token(db, live_user):
|
||||
"""Only a refresh token can be refreshed"""
|
||||
access_token, _ = issue_tokens(db, live_user.id, live_user.organization_id)
|
||||
db.commit()
|
||||
|
||||
assert await refresh_access_token(access_token, db) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_access_token_migrates_a_legacy_token(db, live_user):
|
||||
"""A refresh token from before sessions still works, and gets a session"""
|
||||
legacy = create_refresh_token({"sub": str(live_user.id), "org": str(live_user.organization_id)})
|
||||
|
||||
result = await refresh_access_token(legacy, db)
|
||||
|
||||
sid = verify_token(result)["sid"]
|
||||
session = db.query(AuthSession).one()
|
||||
assert str(session.id) == sid
|
||||
assert session.user_id == live_user.id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_success(socket_db, live_user):
|
||||
"""Test successful socket authentication"""
|
||||
access_token, _ = issue_tokens(socket_db, live_user.id, live_user.organization_id)
|
||||
socket_db.commit()
|
||||
|
||||
result = await authenticate_socket("test_sid", _socket_environ(access_token=access_token))
|
||||
|
||||
assert result == (access_token, str(live_user.id), str(live_user.organization_id))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_with_refresh(socket_db, live_user):
|
||||
"""Test socket authentication with token refresh"""
|
||||
sid = "test_sid"
|
||||
_, refresh_token = issue_tokens(socket_db, live_user.id, live_user.organization_id)
|
||||
socket_db.commit()
|
||||
|
||||
result_token, result_user_id, result_org_id = await authenticate_socket(
|
||||
sid, _socket_environ(refresh_token=refresh_token))
|
||||
|
||||
assert verify_token(result_token)["type"] == "access"
|
||||
assert result_user_id == str(live_user.id)
|
||||
assert result_org_id == str(live_user.organization_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_rejects_a_revoked_session(socket_db, live_user):
|
||||
"""A token copied before sign-out cannot open a socket"""
|
||||
access_token, _ = issue_tokens(socket_db, live_user.id, live_user.organization_id)
|
||||
revoke_session(socket_db, verify_token(access_token)["sid"])
|
||||
socket_db.commit()
|
||||
|
||||
result = await authenticate_socket("test_sid", _socket_environ(access_token=access_token))
|
||||
|
||||
assert result == (None, None, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_rejects_a_refresh_token(socket_db, live_user):
|
||||
"""A refresh token in the access_token cookie is the wrong type"""
|
||||
_, refresh_token = issue_tokens(socket_db, live_user.id, live_user.organization_id)
|
||||
socket_db.commit()
|
||||
|
||||
result = await authenticate_socket("test_sid", _socket_environ(access_token=refresh_token))
|
||||
|
||||
assert result == (None, None, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_rejects_an_inactive_user(socket_db, live_user):
|
||||
"""A deactivated user's still-unexpired token cannot open a socket"""
|
||||
access_token, _ = issue_tokens(socket_db, live_user.id, live_user.organization_id)
|
||||
live_user.is_active = False
|
||||
socket_db.commit()
|
||||
|
||||
result = await authenticate_socket("test_sid", _socket_environ(access_token=access_token))
|
||||
|
||||
assert result == (None, None, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_rejects_a_deleted_user(socket_db):
|
||||
"""A legacy token for a user who no longer exists cannot open a socket"""
|
||||
access_token = create_access_token({"sub": str(uuid4()), "org": str(uuid4())})
|
||||
|
||||
result = await authenticate_socket("test_sid", _socket_environ(access_token=access_token))
|
||||
|
||||
assert result == (None, None, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_conversation_token_success(mock_db, mock_widget):
|
||||
"""Test successful widget socket authentication"""
|
||||
# Setup
|
||||
sid = "test_sid"
|
||||
conversation_token = "valid_conversation_token"
|
||||
customer_id = str(uuid4())
|
||||
|
||||
# Create auth data
|
||||
auth = {'conversation_token': conversation_token}
|
||||
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_widget
|
||||
|
||||
# No merged-into pointer for this customer: it resolves to itself.
|
||||
with patch('app.core.auth_utils.verify_conversation_token') as mock_verify, \
|
||||
patch('app.core.auth_utils.SessionLocal') as mock_session_local, \
|
||||
patch('app.repositories.customer.CustomerRepository.resolve_merged',
|
||||
return_value=MagicMock(id=customer_id)):
|
||||
mock_verify.return_value = {
|
||||
"widget_id": str(mock_widget.id),
|
||||
"sub": customer_id,
|
||||
"type": "conversation"
|
||||
}
|
||||
mock_session_local.return_value.__enter__.return_value = mock_db
|
||||
|
||||
# Execute
|
||||
result_widget_id, result_org_id, result_customer_id, result_token = await authenticate_socket_conversation_token(sid, auth)
|
||||
|
||||
# Assert
|
||||
assert result_widget_id == str(mock_widget.id)
|
||||
assert result_org_id == mock_widget.organization_id
|
||||
assert result_customer_id == customer_id
|
||||
assert result_token == conversation_token
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_conversation_token_invalid(mock_db):
|
||||
"""Test widget socket authentication with invalid token"""
|
||||
# Setup
|
||||
sid = "test_sid"
|
||||
conversation_token = "invalid_token"
|
||||
|
||||
# Create auth data
|
||||
auth = {'conversation_token': conversation_token}
|
||||
|
||||
with patch('app.core.auth_utils.verify_conversation_token') as mock_verify:
|
||||
mock_verify.return_value = None
|
||||
|
||||
# Execute
|
||||
result = await authenticate_socket_conversation_token(sid, auth)
|
||||
|
||||
# Assert
|
||||
assert result == (None, None, None, None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_conversation_token_missing_widget(mock_db):
|
||||
"""Test widget socket authentication with missing widget"""
|
||||
# Setup
|
||||
sid = "test_sid"
|
||||
conversation_token = "valid_conversation_token"
|
||||
customer_id = str(uuid4())
|
||||
widget_id = str(uuid4())
|
||||
|
||||
# Create auth data
|
||||
auth = {'conversation_token': conversation_token}
|
||||
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = None
|
||||
|
||||
with patch('app.core.auth_utils.verify_conversation_token') as mock_verify, \
|
||||
patch('app.core.auth_utils.get_db') as mock_get_db:
|
||||
mock_verify.return_value = {
|
||||
"widget_id": widget_id,
|
||||
"sub": customer_id,
|
||||
"type": "conversation"
|
||||
}
|
||||
mock_get_db.return_value.__next__.return_value = mock_db
|
||||
|
||||
# Execute
|
||||
result = await authenticate_socket_conversation_token(sid, auth)
|
||||
|
||||
# Assert
|
||||
assert result == (None, None, None, None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_conversation_token_wrong_type(mock_db, mock_widget):
|
||||
"""Test widget socket authentication with wrong token type"""
|
||||
# Setup
|
||||
sid = "test_sid"
|
||||
conversation_token = "wrong_type_token"
|
||||
customer_id = str(uuid4())
|
||||
|
||||
# Create auth data
|
||||
auth = {'conversation_token': conversation_token}
|
||||
|
||||
mock_db.query.return_value.filter.return_value.first.return_value = mock_widget
|
||||
|
||||
with patch('app.core.auth_utils.verify_conversation_token') as mock_verify, \
|
||||
patch('app.core.auth_utils.get_db') as mock_get_db:
|
||||
mock_verify.return_value = {
|
||||
"widget_id": str(mock_widget.id),
|
||||
"sub": customer_id,
|
||||
"type": "access" # Wrong token type
|
||||
}
|
||||
mock_get_db.return_value.__next__.return_value = mock_db
|
||||
|
||||
# Execute
|
||||
result = await authenticate_socket_conversation_token(sid, auth)
|
||||
|
||||
# Assert
|
||||
assert result == (None, None, None, None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_no_cookies(mock_db):
|
||||
"""Test socket authentication with no cookies"""
|
||||
# Setup
|
||||
sid = "test_sid"
|
||||
environ = {}
|
||||
|
||||
# Execute
|
||||
result_token, result_user_id, result_org_id = await authenticate_socket(sid, environ)
|
||||
|
||||
# Assert
|
||||
assert result_token is None
|
||||
assert result_user_id is None
|
||||
assert result_org_id is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_socket_invalid_token_format(mock_db):
|
||||
"""Test socket authentication with invalid token format"""
|
||||
# Setup
|
||||
sid = "test_sid"
|
||||
access_token = "invalid_format_token"
|
||||
|
||||
# Create mock environ with cookies
|
||||
environ = _socket_environ(access_token=access_token)
|
||||
|
||||
with patch('app.core.auth_utils.live_payload') as mock_verify:
|
||||
mock_verify.side_effect = Exception("Invalid token format")
|
||||
|
||||
# Execute
|
||||
result_token, result_user_id, result_org_id = await authenticate_socket(sid, environ)
|
||||
|
||||
# Assert
|
||||
assert result_token is None
|
||||
assert result_user_id is None
|
||||
assert result_org_id is None
|
||||
|
||||
|
||||
def test_widget_socket_identity_reads_the_connected_session():
|
||||
"""A live connection's identity comes from what connect resolved, not from
|
||||
re-verifying the token on every event (#315)."""
|
||||
session = {
|
||||
"widget_id": "widget-1",
|
||||
"org_id": "org-1",
|
||||
"customer_id": "customer-1",
|
||||
"conversation_token": "long-since-expired",
|
||||
}
|
||||
|
||||
assert widget_socket_identity(session) == ("widget-1", "org-1", "customer-1")
|
||||
|
||||
|
||||
def test_widget_socket_identity_allows_an_anonymous_customer():
|
||||
"""A visitor who has not been identified still has a widget and an org."""
|
||||
session = {"widget_id": "widget-1", "org_id": "org-1"}
|
||||
|
||||
assert widget_socket_identity(session) == ("widget-1", "org-1", None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("session", [
|
||||
None,
|
||||
{},
|
||||
{"widget_id": "widget-1"},
|
||||
{"org_id": "org-1"},
|
||||
])
|
||||
def test_widget_socket_identity_rejects_an_unauthenticated_socket(session):
|
||||
"""No session means connect never authenticated this sid."""
|
||||
assert widget_socket_identity(session) == (None, None, None)
|
||||
@@ -0,0 +1,132 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from app.core.config import (
|
||||
Settings,
|
||||
check_secret_configuration,
|
||||
verify_secret_configuration,
|
||||
)
|
||||
|
||||
# The key that used to ship as the ENCRYPTION_KEY default and in .env.example. It
|
||||
# is public, so a deployment still using it must be flagged.
|
||||
LEGACY_PUBLIC_KEY = "RFQ4SzhyRTVYdGtsLUxsc25SaDB0QlZpbTdQRmlVRlpsZUlCaFRlU2Vxbz0="
|
||||
GOOD = {
|
||||
"SECRET_KEY": "a-real-secret",
|
||||
"CONVERSATION_SECRET_KEY": "another-real-secret",
|
||||
"ENCRYPTION_KEY": "a-real-key",
|
||||
}
|
||||
|
||||
|
||||
def _config(**overrides) -> Settings:
|
||||
return Settings(ENVIRONMENT="production", **{**GOOD, **overrides})
|
||||
|
||||
|
||||
def test_configured_secrets_are_not_flagged():
|
||||
assert check_secret_configuration(_config()) == []
|
||||
|
||||
|
||||
def test_unset_encryption_key_is_flagged():
|
||||
"""ENCRYPTION_KEY no longer carries a default, so unset is the common case."""
|
||||
assert check_secret_configuration(_config(ENCRYPTION_KEY="")) == ["ENCRYPTION_KEY"]
|
||||
|
||||
|
||||
def test_legacy_public_encryption_key_is_still_flagged():
|
||||
"""Removing the default does not help deployments that already copied it."""
|
||||
assert check_secret_configuration(
|
||||
_config(ENCRYPTION_KEY=LEGACY_PUBLIC_KEY)
|
||||
) == ["ENCRYPTION_KEY"]
|
||||
|
||||
|
||||
def test_env_example_placeholder_is_flagged():
|
||||
assert check_secret_configuration(
|
||||
_config(ENCRYPTION_KEY="your_fernet_encryption_key_here")
|
||||
) == ["ENCRYPTION_KEY"]
|
||||
|
||||
|
||||
def test_default_jwt_secret_is_flagged():
|
||||
assert check_secret_configuration(
|
||||
_config(SECRET_KEY="your-secret-key")
|
||||
) == ["SECRET_KEY"]
|
||||
|
||||
|
||||
def test_development_is_exempt():
|
||||
"""Local development runs on the defaults on purpose."""
|
||||
config = Settings(ENVIRONMENT="development", SECRET_KEY="your-secret-key",
|
||||
ENCRYPTION_KEY="")
|
||||
assert check_secret_configuration(config) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("environment", ["test", "testing", "Development"])
|
||||
def test_throwaway_environments_are_exempt(environment):
|
||||
"""The audit and app.core.encryption must agree on what a real deployment is;
|
||||
they used to disagree on casing and on whether test runs counted."""
|
||||
config = Settings(ENVIRONMENT=environment, SECRET_KEY="your-secret-key",
|
||||
ENCRYPTION_KEY="")
|
||||
assert check_secret_configuration(config) == []
|
||||
|
||||
|
||||
def test_startup_passes_with_real_secrets():
|
||||
verify_secret_configuration(_config())
|
||||
|
||||
|
||||
def test_startup_fails_on_default_jwt_secret():
|
||||
with pytest.raises(RuntimeError, match="Refusing to start") as exc:
|
||||
verify_secret_configuration(_config(SECRET_KEY="your-secret-key"))
|
||||
# The env var, not the setting name, plus a command that produces a value.
|
||||
assert "JWT_SECRET_KEY=$(openssl rand -hex 32)" in str(exc.value)
|
||||
|
||||
|
||||
def test_startup_fails_on_public_encryption_key():
|
||||
with pytest.raises(RuntimeError, match="Refusing to start") as exc:
|
||||
verify_secret_configuration(_config(ENCRYPTION_KEY=LEGACY_PUBLIC_KEY))
|
||||
assert "Fernet.generate_key()" in str(exc.value)
|
||||
|
||||
|
||||
def test_startup_error_names_every_offender():
|
||||
with pytest.raises(RuntimeError) as exc:
|
||||
verify_secret_configuration(
|
||||
Settings(ENVIRONMENT="production", SECRET_KEY="",
|
||||
CONVERSATION_SECRET_KEY="", ENCRYPTION_KEY="")
|
||||
)
|
||||
message = str(exc.value)
|
||||
for name in ("SECRET_KEY", "CONVERSATION_SECRET_KEY", "ENCRYPTION_KEY"):
|
||||
assert name in message
|
||||
|
||||
|
||||
def test_development_still_starts_on_defaults():
|
||||
"""The escape hatch is for production; development must stay zero-config."""
|
||||
verify_secret_configuration(
|
||||
Settings(ENVIRONMENT="development", SECRET_KEY="your-secret-key",
|
||||
ENCRYPTION_KEY="")
|
||||
)
|
||||
|
||||
|
||||
def test_escape_hatch_allows_only_the_encryption_key():
|
||||
"""An instance whose data is encrypted under the public key can still boot."""
|
||||
verify_secret_configuration(
|
||||
_config(ENCRYPTION_KEY=LEGACY_PUBLIC_KEY, ALLOW_INSECURE_ENCRYPTION_KEY=True)
|
||||
)
|
||||
|
||||
|
||||
def test_escape_hatch_does_not_excuse_the_jwt_secrets():
|
||||
"""Rotating a JWT secret only ends sessions, so it has no reason to be exempt."""
|
||||
with pytest.raises(RuntimeError, match="SECRET_KEY"):
|
||||
verify_secret_configuration(
|
||||
_config(SECRET_KEY="your-secret-key", ENCRYPTION_KEY=LEGACY_PUBLIC_KEY,
|
||||
ALLOW_INSECURE_ENCRYPTION_KEY=True)
|
||||
)
|
||||
@@ -0,0 +1,292 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
|
||||
Only our own surfaces may use the dashboard's sign-in cookies cross-origin.
|
||||
|
||||
An org's domain is whatever was typed at signup, unverified, and every org
|
||||
domain is in the CORS list. If those origins were also allowed credentials, a
|
||||
page on any signed-up domain could read a signed-in dashboard user's account.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import app.main # noqa: F401 - the CORS middleware is added on import
|
||||
from app.core import auth_utils, cors
|
||||
from app.core.application import app as api
|
||||
|
||||
DASHBOARD = "https://app.example.com"
|
||||
OPERATOR_LISTED = "https://admin.example.com"
|
||||
ORG_SITE = "https://signed-up-domain.example"
|
||||
HELP_CENTER = "https://acme.help.example.com"
|
||||
API_HOST = "api.example.com"
|
||||
|
||||
|
||||
def configure(monkeypatch, frontend_url=DASHBOARD, cors_origins=(OPERATOR_LISTED, DASHBOARD + "/")):
|
||||
monkeypatch.setattr(cors.settings, "FRONTEND_URL", frontend_url)
|
||||
monkeypatch.setattr(cors, "BASE_CORS_ORIGINS", list(cors_origins))
|
||||
cors.first_party_origins.cache_clear()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def first_party(monkeypatch):
|
||||
configure(monkeypatch)
|
||||
yield
|
||||
# Don't leave this module's origins cached for other tests
|
||||
cors.first_party_origins.cache_clear()
|
||||
|
||||
|
||||
def cors_app(allow_origins=(DASHBOARD, OPERATOR_LISTED, ORG_SITE)):
|
||||
app = FastAPI()
|
||||
app.add_middleware(
|
||||
cors.FirstPartyCredentialsCORSMiddleware,
|
||||
credentialed_origins=cors.first_party_origins(),
|
||||
allow_origins=list(allow_origins),
|
||||
allow_origin_regex=r"^https://[a-z0-9-]+\.help\.example\.com$",
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
@app.get("/me")
|
||||
def me():
|
||||
return {"who": "a signed-in user"}
|
||||
|
||||
@app.get("/refused")
|
||||
def refused():
|
||||
raise HTTPException(status_code=401, detail="Not signed in")
|
||||
|
||||
return app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
return TestClient(cors_app())
|
||||
|
||||
|
||||
def preflight(client, origin):
|
||||
return client.options("/me", headers={"Origin": origin, "Access-Control-Request-Method": "GET"})
|
||||
|
||||
|
||||
def cors_layers(app):
|
||||
return [m for m in app.user_middleware if isinstance(m.cls, type) and issubclass(m.cls, cors.CORSMiddleware)]
|
||||
|
||||
|
||||
class TestFirstPartyOrigins:
|
||||
def test_the_dashboard_and_operator_listed_origins(self):
|
||||
assert cors.first_party_origins() == {DASHBOARD, OPERATOR_LISTED}
|
||||
|
||||
@pytest.mark.parametrize("frontend_url, expected", [
|
||||
("https://App.Example.com:443/dashboard/", "https://app.example.com"),
|
||||
("http://localhost:5173/", "http://localhost:5173"),
|
||||
("http://localhost:80", "http://localhost"),
|
||||
("https://app.example.com:8443", "https://app.example.com:8443"),
|
||||
])
|
||||
def test_frontend_url_is_read_the_way_browsers_send_origin(self, monkeypatch, frontend_url, expected):
|
||||
configure(monkeypatch, frontend_url=frontend_url)
|
||||
assert expected in cors.first_party_origins()
|
||||
|
||||
def test_a_frontend_url_without_a_scheme_is_reported(self, monkeypatch, caplog):
|
||||
configure(monkeypatch, frontend_url="app.example.com")
|
||||
with caplog.at_level(logging.ERROR):
|
||||
assert cors.first_party_origins() == {OPERATOR_LISTED, DASHBOARD}
|
||||
assert "FRONTEND_URL" in caplog.text
|
||||
|
||||
def test_a_wildcard_in_cors_origins_trusts_no_one(self, monkeypatch):
|
||||
configure(monkeypatch, cors_origins=["*"])
|
||||
assert cors.first_party_origins() == {DASHBOARD}
|
||||
assert not cors.is_first_party_origin(ORG_SITE)
|
||||
|
||||
@pytest.mark.parametrize("origin, expected", [
|
||||
(DASHBOARD, True),
|
||||
(DASHBOARD.upper(), True),
|
||||
(DASHBOARD + "/", True),
|
||||
(DASHBOARD + ":443", True),
|
||||
(OPERATOR_LISTED, True),
|
||||
("http://app.example.com", False),
|
||||
("https://app.example.com:8443", False),
|
||||
(ORG_SITE, False),
|
||||
(HELP_CENTER, False),
|
||||
("null", False),
|
||||
("", False),
|
||||
(None, False),
|
||||
])
|
||||
def test_is_first_party(self, origin, expected):
|
||||
assert cors.is_first_party_origin(origin) is expected
|
||||
|
||||
|
||||
class TestCredentials:
|
||||
@pytest.mark.parametrize("origin", [DASHBOARD, OPERATOR_LISTED])
|
||||
def test_our_own_origins_may_send_credentials(self, client, origin):
|
||||
for response in (client.get("/me", headers={"Origin": origin}), preflight(client, origin)):
|
||||
assert response.headers["access-control-allow-origin"] == origin
|
||||
assert response.headers["access-control-allow-credentials"] == "true"
|
||||
|
||||
@pytest.mark.parametrize("origin", [ORG_SITE, HELP_CENTER])
|
||||
def test_an_org_site_may_call_the_api_but_never_with_cookies(self, client, origin):
|
||||
for response in (client.get("/me", headers={"Origin": origin}), preflight(client, origin)):
|
||||
assert response.headers["access-control-allow-origin"] == origin
|
||||
assert "access-control-allow-credentials" not in response.headers
|
||||
|
||||
def test_even_when_it_sends_cookies(self, client):
|
||||
response = client.get("/me", headers={"Origin": ORG_SITE, "Cookie": "access_token=x"})
|
||||
assert "access-control-allow-credentials" not in response.headers
|
||||
|
||||
def test_error_responses_too(self, client):
|
||||
org = client.get("/refused", headers={"Origin": ORG_SITE})
|
||||
assert org.status_code == 401
|
||||
assert org.headers["access-control-allow-origin"] == ORG_SITE
|
||||
assert "access-control-allow-credentials" not in org.headers
|
||||
assert client.get("/refused", headers={"Origin": DASHBOARD}).headers["access-control-allow-credentials"] == "true"
|
||||
|
||||
def test_a_wildcard_allow_list_still_keeps_cookies_first_party(self):
|
||||
client = TestClient(cors_app(allow_origins=["*"]))
|
||||
org = client.get("/me", headers={"Origin": ORG_SITE, "Cookie": "access_token=x"})
|
||||
assert "access-control-allow-credentials" not in org.headers
|
||||
dashboard = client.get("/me", headers={"Origin": DASHBOARD, "Cookie": "access_token=x"})
|
||||
assert dashboard.headers["access-control-allow-origin"] == DASHBOARD
|
||||
assert dashboard.headers["access-control-allow-credentials"] == "true"
|
||||
|
||||
def test_an_unknown_site_gets_nothing(self, client):
|
||||
response = client.get("/me", headers={"Origin": "https://evil.example"})
|
||||
assert "access-control-allow-origin" not in response.headers
|
||||
assert "access-control-allow-credentials" not in response.headers
|
||||
assert preflight(client, "https://evil.example").status_code == 400
|
||||
|
||||
|
||||
class TestRefreshKeepsTheSplit:
|
||||
def test_update_local_cors_keeps_first_party_credentials(self, client):
|
||||
app = client.app
|
||||
cors.update_local_cors(app, [DASHBOARD, ORG_SITE])
|
||||
|
||||
layers = cors_layers(app)
|
||||
assert len(layers) == 1
|
||||
assert layers[0].cls is cors.FirstPartyCredentialsCORSMiddleware
|
||||
assert "access-control-allow-credentials" not in client.get("/me", headers={"Origin": ORG_SITE}).headers
|
||||
assert client.get("/me", headers={"Origin": DASHBOARD}).headers["access-control-allow-credentials"] == "true"
|
||||
|
||||
def test_the_apis_own_origin_is_always_allowed(self, monkeypatch):
|
||||
# The iframe snippet and path-mode help centres are pages on the API
|
||||
monkeypatch.setattr(cors.settings, "BACKEND_URL", "https://api.example.com/")
|
||||
assert {"https://api.example.com", ORG_SITE} <= set(cors.with_configured_origins([ORG_SITE]))
|
||||
assert not cors.is_first_party_origin("https://api.example.com")
|
||||
|
||||
def test_a_cached_list_from_an_older_config_never_drops_the_dashboard(self, client):
|
||||
# Redis keeps the list across restarts; it may predate the current env
|
||||
cors.update_local_cors(client.app, [ORG_SITE])
|
||||
|
||||
response = client.get("/me", headers={"Origin": DASHBOARD})
|
||||
assert response.headers["access-control-allow-origin"] == DASHBOARD
|
||||
assert response.headers["access-control-allow-credentials"] == "true"
|
||||
|
||||
def test_the_same_origins_in_another_order_do_not_rebuild(self, client):
|
||||
origins = cors.with_configured_origins([ORG_SITE, "https://other.example"])
|
||||
cors.update_local_cors(client.app, origins)
|
||||
client.get("/me")
|
||||
built = client.app.middleware_stack
|
||||
|
||||
cors.update_local_cors(client.app, list(reversed(origins)))
|
||||
|
||||
assert client.app.middleware_stack is built
|
||||
|
||||
def test_a_rebuilt_middleware_is_the_first_party_one(self):
|
||||
bare = FastAPI()
|
||||
cors.update_local_cors(bare, [ORG_SITE])
|
||||
assert bare.user_middleware[0].cls is cors.FirstPartyCredentialsCORSMiddleware
|
||||
|
||||
|
||||
def test_the_app_is_wired_with_the_first_party_middleware():
|
||||
layers = cors_layers(api)
|
||||
assert len(layers) == 1
|
||||
assert layers[0].cls is cors.FirstPartyCredentialsCORSMiddleware
|
||||
assert layers[0].kwargs["credentialed_origins"]
|
||||
|
||||
|
||||
def handshake(origin=None, query="EIO=4&transport=polling", host=API_HOST, fetch_site=None, cookie="access_token=tok"):
|
||||
environ = {"QUERY_STRING": query, "HTTP_HOST": host, "HTTP_COOKIE": cookie}
|
||||
if origin:
|
||||
environ["HTTP_ORIGIN"] = origin
|
||||
if fetch_site:
|
||||
environ["HTTP_SEC_FETCH_SITE"] = fetch_site
|
||||
return environ
|
||||
|
||||
|
||||
class TestCookieSockets:
|
||||
"""Sockets get no CORS check, and socket.io signs a namespace in with the
|
||||
headers of the first handshake request, so the handshake itself must have
|
||||
come from our own pages."""
|
||||
|
||||
@pytest.mark.parametrize("environ", [
|
||||
handshake(origin=DASHBOARD),
|
||||
handshake(origin=DASHBOARD + ":443"),
|
||||
handshake(origin=OPERATOR_LISTED),
|
||||
handshake(origin=f"https://{API_HOST}"),
|
||||
handshake(origin=f"https://{API_HOST}:443"),
|
||||
handshake(origin="http://127.0.0.1:8000", host="127.0.0.1:8000"),
|
||||
handshake(fetch_site="same-origin"),
|
||||
], ids=["dashboard", "dashboard-default-port", "operator-listed", "the-api-itself", "the-api-default-port",
|
||||
"the-api-on-a-port", "same-origin-no-origin"])
|
||||
def test_our_own_pages_may_use_the_cookies(self, environ):
|
||||
assert auth_utils.cookie_socket_allowed(environ)
|
||||
|
||||
@pytest.mark.parametrize("environ", [
|
||||
handshake(origin=ORG_SITE),
|
||||
handshake(origin=HELP_CENTER),
|
||||
handshake(origin="null"),
|
||||
handshake(),
|
||||
handshake(fetch_site="cross-site"),
|
||||
handshake(fetch_site="same-site"),
|
||||
handshake(query="EIO=4&transport=polling&j=0"),
|
||||
handshake(query="EIO=4&transport=polling&j=0", fetch_site="same-origin"),
|
||||
handshake(origin=DASHBOARD, query="EIO=4&transport=polling&j=0"),
|
||||
handshake(origin=ORG_SITE, host="signed-up-domain.example:8443"),
|
||||
handshake(origin=f"https://{API_HOST}", host=""),
|
||||
], ids=["org-site", "help-centre", "opaque", "no-origin", "cross-site-no-origin", "same-site-no-origin",
|
||||
"jsonp", "jsonp-same-origin", "jsonp-from-dashboard", "org-site-other-port", "no-host"])
|
||||
def test_anything_else_may_not(self, environ):
|
||||
assert not auth_utils.cookie_socket_allowed(environ)
|
||||
|
||||
@pytest.fixture
|
||||
def auth_calls(self, monkeypatch):
|
||||
calls = []
|
||||
|
||||
async def refresh(token, db):
|
||||
calls.append(("refresh", token))
|
||||
return "fresh-access-token"
|
||||
|
||||
monkeypatch.setattr(auth_utils, "refresh_access_token", refresh)
|
||||
monkeypatch.setattr(auth_utils, "live_payload", lambda db, token, kind: calls.append(("verify", token)) or None)
|
||||
return calls
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("environ", [
|
||||
handshake(origin=ORG_SITE, cookie="refresh_token=r; access_token=a"),
|
||||
handshake(cookie="refresh_token=r", query="EIO=4&transport=polling&j=0"),
|
||||
], ids=["org-site", "jsonp"])
|
||||
async def test_a_refused_handshake_never_reads_or_refreshes_the_cookies(self, auth_calls, environ):
|
||||
assert await auth_utils.authenticate_socket("sid-1", environ) == (None, None, None)
|
||||
assert auth_calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_our_dashboard_is_still_authenticated_from_its_cookies(self, auth_calls):
|
||||
await auth_utils.authenticate_socket("sid-1", handshake(origin=DASHBOARD))
|
||||
assert auth_calls == [("verify", "tok")]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_refreshed_token_stays_on_the_server(self, auth_calls):
|
||||
await auth_utils.authenticate_socket("sid-1", handshake(origin=DASHBOARD, cookie="refresh_token=r"))
|
||||
assert auth_calls == [("refresh", "r"), ("verify", "fresh-access-token")]
|
||||
@@ -0,0 +1,156 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
|
||||
Refreshing the allowed origins must update the one CORS middleware in place.
|
||||
It used to prepend a new layer on every refresh, and the stale inner layers
|
||||
kept answering for origins that had been removed.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.core import cors
|
||||
from app.core.socketio import sio
|
||||
|
||||
KEPT = "https://kept.example.com"
|
||||
REMOVED = "https://removed.example.com"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def restore_socketio_origins(monkeypatch):
|
||||
# update_local_cors also rewrites the shared Socket.IO server's list
|
||||
monkeypatch.setattr(sio.eio, "cors_allowed_origins", list(sio.eio.cors_allowed_origins or []))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app():
|
||||
app = FastAPI()
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=[KEPT, REMOVED],
|
||||
allow_origin_regex=cors.get_cors_origin_regex(),
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
@app.get("/ping")
|
||||
def ping():
|
||||
return {"ok": True}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def cors_layers(app):
|
||||
return [m for m in app.user_middleware if isinstance(m.cls, type) and issubclass(m.cls, CORSMiddleware)]
|
||||
|
||||
|
||||
def get_from(client, origin):
|
||||
return client.get("/ping", headers={"Origin": origin})
|
||||
|
||||
|
||||
def test_refresh_updates_the_existing_middleware(app):
|
||||
assert cors.update_local_cors(app, [KEPT]) is True
|
||||
|
||||
layers = cors_layers(app)
|
||||
assert len(layers) == 1
|
||||
assert layers[0].kwargs["allow_origins"] == cors.with_configured_origins([KEPT])
|
||||
assert layers[0].kwargs["allow_origin_regex"] == cors.get_cors_origin_regex()
|
||||
|
||||
|
||||
def test_repeated_refreshes_never_stack_layers(app):
|
||||
client = TestClient(app)
|
||||
get_from(client, KEPT) # the stack is built once the app has served a request
|
||||
|
||||
for _ in range(3):
|
||||
cors.update_local_cors(app, [KEPT])
|
||||
get_from(client, KEPT)
|
||||
|
||||
assert len(cors_layers(app)) == 1
|
||||
|
||||
|
||||
def test_removed_origin_is_no_longer_allowed(app):
|
||||
client = TestClient(app)
|
||||
assert get_from(client, REMOVED).headers.get("access-control-allow-origin") == REMOVED
|
||||
|
||||
cors.update_local_cors(app, [KEPT])
|
||||
|
||||
assert "access-control-allow-origin" not in get_from(client, REMOVED).headers
|
||||
assert get_from(client, KEPT).headers.get("access-control-allow-origin") == KEPT
|
||||
|
||||
|
||||
def test_removed_origin_preflight_is_refused(app):
|
||||
client = TestClient(app)
|
||||
cors.update_local_cors(app, [KEPT])
|
||||
|
||||
response = client.options(
|
||||
"/ping",
|
||||
headers={"Origin": REMOVED, "Access-Control-Request-Method": "GET"},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "access-control-allow-origin" not in response.headers
|
||||
|
||||
|
||||
def test_vary_header_is_not_duplicated(app):
|
||||
client = TestClient(app)
|
||||
for _ in range(3):
|
||||
cors.update_local_cors(app, [KEPT])
|
||||
|
||||
vary = get_from(client, KEPT).headers.get("vary", "")
|
||||
assert [v.strip() for v in vary.split(",")].count("Origin") == 1
|
||||
|
||||
|
||||
def test_help_center_regex_survives_a_refresh(app, monkeypatch):
|
||||
monkeypatch.setattr(cors.settings, "HELP_CENTER_BASE_DOMAIN", "help.example.com")
|
||||
cors.get_cors_origin_regex.cache_clear()
|
||||
try:
|
||||
app.user_middleware[0].kwargs["allow_origin_regex"] = cors.get_cors_origin_regex()
|
||||
client = TestClient(app)
|
||||
cors.update_local_cors(app, [KEPT])
|
||||
|
||||
origin = "https://acme.help.example.com"
|
||||
assert get_from(client, origin).headers.get("access-control-allow-origin") == origin
|
||||
finally:
|
||||
cors.get_cors_origin_regex.cache_clear()
|
||||
|
||||
|
||||
def test_unchanged_origins_skip_the_rebuild(app):
|
||||
client = TestClient(app)
|
||||
cors.update_local_cors(app, [KEPT])
|
||||
get_from(client, KEPT)
|
||||
built = app.middleware_stack
|
||||
|
||||
cors.update_local_cors(app, [KEPT])
|
||||
|
||||
assert app.middleware_stack is built
|
||||
|
||||
|
||||
def test_socketio_origins_follow_the_refresh(app):
|
||||
cors.update_local_cors(app, [KEPT])
|
||||
|
||||
assert sio.eio.cors_allowed_origins == cors.with_configured_origins([KEPT])
|
||||
|
||||
|
||||
def test_app_without_cors_middleware_gets_one():
|
||||
bare = FastAPI()
|
||||
|
||||
assert cors.update_local_cors(bare, [KEPT]) is True
|
||||
|
||||
layers = cors_layers(bare)
|
||||
assert len(layers) == 1
|
||||
assert layers[0].kwargs["allow_origins"] == cors.with_configured_origins([KEPT])
|
||||
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
|
||||
CORS admits each org's site with and without www., however the domain was
|
||||
typed at signup - the same rule install detection uses.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from app.core import cors
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def origins_for(monkeypatch):
|
||||
def build(*domains):
|
||||
monkeypatch.setattr("app.repositories.organization.OrganizationRepository.get_active_organizations",
|
||||
lambda self: [SimpleNamespace(domain=d) for d in domains])
|
||||
monkeypatch.setattr("app.repositories.help_center.HelpCenterRepository.list_enabled_slugs", lambda self: [])
|
||||
monkeypatch.setattr("app.repositories.help_center.HelpCenterRepository.list_verified_domains", lambda self: [])
|
||||
cors.get_cors_origins.cache_clear()
|
||||
try:
|
||||
return cors.get_cors_origins()
|
||||
finally:
|
||||
cors.get_cors_origins.cache_clear()
|
||||
return build
|
||||
|
||||
|
||||
@pytest.mark.parametrize("typed", ["acme.com", "www.acme.com", "Acme.com", "https://WWW.Acme.com/"])
|
||||
def test_org_site_admitted_with_and_without_www(origins_for, typed):
|
||||
origins = origins_for(typed)
|
||||
assert {"https://acme.com", "https://www.acme.com", "http://acme.com", "http://www.acme.com"} <= origins
|
||||
assert not any("www.www." in o or "Acme" in o for o in origins)
|
||||
|
||||
|
||||
def test_blank_domain_adds_nothing(origins_for):
|
||||
assert not any(o.endswith("://") or o.endswith("://www.") for o in origins_for("", None))
|
||||
@@ -0,0 +1,104 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
|
||||
The API origin's crawler policy. Search Console reported dozens of
|
||||
/api/v1/... endpoints as "Not found (404)" because this host answered
|
||||
robots.txt with a 404 of its own, so these run against the real wired app
|
||||
rather than a stand-in — the wiring is the part that was missing.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import app.main # noqa: F401 — the robots route and middleware are wired on import
|
||||
from app.core.application import app
|
||||
from app.core.config import settings
|
||||
from app.core.crawler_policy import api_robots_txt
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def client():
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
class TestApiRobotsTxt:
|
||||
def test_served_as_plain_text_not_json(self, client):
|
||||
response = client.get("/robots.txt")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"].startswith("text/plain")
|
||||
|
||||
def test_disallows_everything(self, client):
|
||||
lines = client.get("/robots.txt").text.splitlines()
|
||||
|
||||
assert "User-agent: *" in lines
|
||||
assert "Disallow: /" in lines
|
||||
|
||||
def test_path_mode_keeps_the_self_hosted_help_center_crawlable(self, monkeypatch):
|
||||
"""Path dispatch serves public help-center pages from this same origin,
|
||||
and those are meant to be indexed."""
|
||||
monkeypatch.setattr(settings, "HELP_CENTER_PUBLIC_MODE", "path")
|
||||
|
||||
assert "Allow: /help/" in api_robots_txt().splitlines()
|
||||
|
||||
def test_path_mode_allow_precedes_the_blanket_disallow(self, monkeypatch):
|
||||
"""A crawler that takes the first match instead of the most specific
|
||||
one must not read "Disallow: /" before reaching the carve-out."""
|
||||
monkeypatch.setattr(settings, "HELP_CENTER_PUBLIC_MODE", "path")
|
||||
lines = api_robots_txt().splitlines()
|
||||
|
||||
assert lines.index("Allow: /help/") < lines.index("Disallow: /")
|
||||
|
||||
def test_subdomain_mode_has_nothing_to_carve_out(self, monkeypatch):
|
||||
monkeypatch.setattr(settings, "HELP_CENTER_PUBLIC_MODE", "subdomain")
|
||||
lines = api_robots_txt().splitlines()
|
||||
|
||||
assert "Disallow: /" in lines
|
||||
assert not any(line.startswith("Allow:") for line in lines)
|
||||
|
||||
|
||||
class TestNoIndexHeader:
|
||||
"""robots.txt stops the fetch; this stops a URL discovered elsewhere from
|
||||
being indexed regardless."""
|
||||
|
||||
def test_api_success_response_is_noindex(self, client):
|
||||
probe = Path("uploads") / "agents" / "crawler-policy-probe.txt"
|
||||
probe.parent.mkdir(parents=True, exist_ok=True)
|
||||
probe.write_text("probe")
|
||||
try:
|
||||
response = client.get("/api/v1/uploads/agents/crawler-policy-probe.txt")
|
||||
finally:
|
||||
probe.unlink(missing_ok=True)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["x-robots-tag"] == "noindex"
|
||||
|
||||
def test_api_404_is_noindex(self, client):
|
||||
"""The 404s are exactly what Search Console indexed, so they need it
|
||||
most. Deliberately a path no router claims: the 404 is then raised
|
||||
before any endpoint dependency runs, so this asserts the middleware
|
||||
and not some unrelated route's database access."""
|
||||
response = client.get("/api/v1/no-such-endpoint")
|
||||
|
||||
assert response.status_code == 404
|
||||
assert response.headers["x-robots-tag"] == "noindex"
|
||||
|
||||
def test_non_api_paths_are_left_alone(self, client):
|
||||
response = client.get("/health")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "x-robots-tag" not in response.headers
|
||||
@@ -0,0 +1,47 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from app import database
|
||||
|
||||
|
||||
def test_a_child_starts_with_its_own_pool():
|
||||
parent_pool = database.engine.pool
|
||||
|
||||
database._fresh_pool_after_fork()
|
||||
|
||||
assert database.engine.pool is not parent_pool
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs fork()")
|
||||
def test_a_forked_worker_never_reuses_the_parents_pool():
|
||||
parent_pool_id = id(database.engine.pool)
|
||||
read_end, write_end = os.pipe()
|
||||
pid = os.fork()
|
||||
if pid == 0: # the child: report whether it got a fresh pool, then leave at once
|
||||
os.close(read_end)
|
||||
os.write(write_end, b"fresh" if id(database.engine.pool) != parent_pool_id else b"shared")
|
||||
os._exit(0)
|
||||
os.close(write_end)
|
||||
answer = os.read(read_end, 16)
|
||||
os.close(read_end)
|
||||
os.waitpid(pid, 0)
|
||||
|
||||
assert answer == b"fresh"
|
||||
assert id(database.engine.pool) == parent_pool_id
|
||||
@@ -0,0 +1,53 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from app.core import edition
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clean_meta_path():
|
||||
saved = list(sys.meta_path)
|
||||
yield
|
||||
sys.meta_path[:] = saved
|
||||
|
||||
|
||||
def _blockers():
|
||||
return [f for f in sys.meta_path if isinstance(f, edition._BlockEnterprise)]
|
||||
|
||||
|
||||
def test_apply_without_flag_installs_nothing(monkeypatch, clean_meta_path):
|
||||
monkeypatch.delenv("NO_ENTERPRISE", raising=False)
|
||||
sys.meta_path[:] = [f for f in sys.meta_path if not isinstance(f, edition._BlockEnterprise)]
|
||||
|
||||
edition.apply()
|
||||
|
||||
assert _blockers() == []
|
||||
|
||||
|
||||
def test_apply_with_flag_blocks_enterprise_imports_once(monkeypatch, clean_meta_path):
|
||||
monkeypatch.setenv("NO_ENTERPRISE", "1")
|
||||
|
||||
edition.apply()
|
||||
edition.apply()
|
||||
|
||||
assert len(_blockers()) == 1
|
||||
with pytest.raises(ModuleNotFoundError):
|
||||
_blockers()[0].find_spec("app.enterprise.models")
|
||||
assert _blockers()[0].find_spec("app.models") is None
|
||||
@@ -0,0 +1,246 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from cryptography.fernet import Fernet
|
||||
from sqlalchemy import text
|
||||
|
||||
from app.core.config import settings
|
||||
from app.core.encryption import (
|
||||
CURRENT_KEY_ID,
|
||||
ENCRYPTED_PREFIX,
|
||||
ENCRYPTION_KEY_ENV,
|
||||
JSON_MARKER,
|
||||
SEPARATOR,
|
||||
EncryptedText,
|
||||
decrypt_json,
|
||||
decrypt_value,
|
||||
encrypt_json,
|
||||
encrypt_value,
|
||||
is_encrypted,
|
||||
load_key,
|
||||
)
|
||||
from app.database import Base
|
||||
from app.models.chat_history import ChatHistory
|
||||
from app.models.session_to_agent import SessionToAgent, SessionStatus
|
||||
|
||||
# Where an installed extension's models live (see app.channels.app_hooks)
|
||||
EXTENSION_PACKAGE = "app.enterprise."
|
||||
|
||||
|
||||
def test_round_trip():
|
||||
token = encrypt_value("hello customer")
|
||||
assert token.startswith(f"{ENCRYPTED_PREFIX}{SEPARATOR}{CURRENT_KEY_ID}{SEPARATOR}")
|
||||
assert "hello customer" not in token
|
||||
assert decrypt_value(token) == "hello customer"
|
||||
|
||||
|
||||
def test_encrypt_is_idempotent():
|
||||
token = encrypt_value("hello")
|
||||
assert encrypt_value(token) == token
|
||||
|
||||
|
||||
def test_plaintext_passes_through():
|
||||
"""Rows the backfill has not reached yet must keep reading."""
|
||||
assert decrypt_value("not yet encrypted") == "not yet encrypted"
|
||||
|
||||
|
||||
def test_none_and_empty():
|
||||
assert encrypt_value(None) is None
|
||||
assert decrypt_value(None) is None
|
||||
assert decrypt_value("") == ""
|
||||
assert decrypt_value(encrypt_value("")) == ""
|
||||
|
||||
|
||||
def test_unknown_key_id_returns_none():
|
||||
token = encrypt_value("secret").replace(CURRENT_KEY_ID, "v99", 1)
|
||||
assert decrypt_value(token) is None
|
||||
|
||||
|
||||
def test_corrupt_token_returns_none():
|
||||
assert decrypt_value(f"{ENCRYPTED_PREFIX}{SEPARATOR}{CURRENT_KEY_ID}{SEPARATOR}nope") is None
|
||||
|
||||
|
||||
def test_malformed_ciphertext_returns_none():
|
||||
assert decrypt_value(f"{ENCRYPTED_PREFIX}{SEPARATOR}{CURRENT_KEY_ID}") is None
|
||||
|
||||
|
||||
def test_is_encrypted():
|
||||
assert is_encrypted(encrypt_value("x"))
|
||||
assert not is_encrypted("plain")
|
||||
assert not is_encrypted(None)
|
||||
|
||||
|
||||
def test_json_round_trip():
|
||||
memory = {"runs": [{"content": "my secret transcript"}]}
|
||||
blob = encrypt_json(memory)
|
||||
assert list(blob) == [JSON_MARKER]
|
||||
assert "secret" not in str(blob)
|
||||
assert decrypt_json(blob) == memory
|
||||
|
||||
|
||||
def test_json_plaintext_passes_through():
|
||||
legacy = {"runs": [{"content": "written before encryption"}]}
|
||||
assert decrypt_json(legacy) == legacy
|
||||
assert decrypt_json(None) is None
|
||||
assert encrypt_json(None) is None
|
||||
|
||||
|
||||
def test_json_encrypt_is_idempotent():
|
||||
blob = encrypt_json({"a": 1})
|
||||
assert encrypt_json(blob) == blob
|
||||
|
||||
|
||||
def test_undecryptable_json_returns_none():
|
||||
assert decrypt_json({JSON_MARKER: f"{ENCRYPTED_PREFIX}{SEPARATOR}{CURRENT_KEY_ID}{SEPARATOR}x"}) is None
|
||||
|
||||
|
||||
def test_backfill_script_covers_every_encrypted_column():
|
||||
"""The backfill lives outside the app tree and cannot import the models, so
|
||||
nothing but this test stops a newly encrypted column from shipping with its
|
||||
existing rows silently left in plaintext."""
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "scripts"))
|
||||
try:
|
||||
from encrypt_chat_at_rest import TARGETS
|
||||
finally:
|
||||
sys.path.pop(0)
|
||||
|
||||
# Only the open-source models: an installed extension (app.enterprise)
|
||||
# owns its tables and any backfill they need.
|
||||
declared = {
|
||||
(mapper.local_table.name, column.name)
|
||||
for mapper in Base.registry.mappers
|
||||
if not mapper.class_.__module__.startswith(EXTENSION_PACKAGE)
|
||||
for column in mapper.local_table.columns
|
||||
if isinstance(column.type, EncryptedText)
|
||||
}
|
||||
covered = {
|
||||
(table, column)
|
||||
for table, _pk, columns, _kind in TARGETS
|
||||
for column in columns
|
||||
}
|
||||
assert declared <= covered, f"not backfilled: {declared - covered}"
|
||||
|
||||
|
||||
def test_chat_message_is_ciphertext_in_the_database(db, test_organization_id):
|
||||
"""The ORM must round-trip, and the stored bytes must not be readable."""
|
||||
message = ChatHistory(
|
||||
organization_id=test_organization_id,
|
||||
message="card number is 4111 1111 1111 1111",
|
||||
message_type="user",
|
||||
)
|
||||
db.add(message)
|
||||
db.commit()
|
||||
|
||||
# Raw SQL: selecting the typed column would decrypt it again on the way out.
|
||||
raw = db.execute(
|
||||
text("SELECT message FROM chat_history WHERE id = :id"), {"id": message.id}
|
||||
).scalar()
|
||||
assert is_encrypted(raw)
|
||||
assert "4111" not in raw
|
||||
|
||||
db.expire_all()
|
||||
assert db.get(ChatHistory, message.id).message == "card number is 4111 1111 1111 1111"
|
||||
|
||||
|
||||
def test_session_summary_is_ciphertext_in_the_database(db, test_organization_id):
|
||||
session = SessionToAgent(
|
||||
session_id=uuid4(),
|
||||
organization_id=test_organization_id,
|
||||
customer_id=uuid4(),
|
||||
status=SessionStatus.OPEN,
|
||||
ticket_summary="customer threatened to churn",
|
||||
transfer_description="needs a human",
|
||||
)
|
||||
db.add(session)
|
||||
db.commit()
|
||||
|
||||
raw = db.execute(
|
||||
text("SELECT ticket_summary FROM session_to_agents "
|
||||
"WHERE ticket_summary IS NOT NULL")
|
||||
).scalar()
|
||||
assert is_encrypted(raw)
|
||||
assert "churn" not in raw
|
||||
|
||||
db.expire_all()
|
||||
reloaded = db.get(SessionToAgent, session.session_id)
|
||||
assert reloaded.ticket_summary == "customer threatened to churn"
|
||||
assert reloaded.transfer_description == "needs a human"
|
||||
|
||||
|
||||
def test_legacy_plaintext_row_still_reads(db, test_organization_id):
|
||||
"""A row written before this change has no prefix and must not break reads."""
|
||||
# Raw SQL on purpose: going through the ORM would encrypt it.
|
||||
db.execute(
|
||||
text("INSERT INTO chat_history (organization_id, message, message_type) "
|
||||
"VALUES (:org, :message, :message_type)"),
|
||||
{"org": str(test_organization_id), "message": "written before encryption",
|
||||
"message_type": "user"},
|
||||
)
|
||||
db.commit()
|
||||
|
||||
stored = db.query(ChatHistory).filter(
|
||||
ChatHistory.message_type == "user"
|
||||
).order_by(ChatHistory.id.desc()).first()
|
||||
assert stored.message == "written before encryption"
|
||||
|
||||
|
||||
def _double_encoded(key: bytes) -> str:
|
||||
"""The wire format ENCRYPTION_KEY uses: a Fernet key, base64-encoded again."""
|
||||
return base64.b64encode(key).decode()
|
||||
|
||||
|
||||
def test_load_key_accepts_a_double_encoded_key(monkeypatch):
|
||||
monkeypatch.setattr(settings, "ENVIRONMENT", "production")
|
||||
key = Fernet.generate_key()
|
||||
monkeypatch.setenv(ENCRYPTION_KEY_ENV, _double_encoded(key))
|
||||
assert load_key() == key
|
||||
|
||||
|
||||
def test_load_key_rejects_base64_that_is_not_a_fernet_key(monkeypatch):
|
||||
"""Encoding the key once instead of twice decodes cleanly but is 32 raw bytes,
|
||||
not Fernet key material — it must fail here, not later inside Fernet()."""
|
||||
monkeypatch.setattr(settings, "ENVIRONMENT", "production")
|
||||
monkeypatch.setenv(ENCRYPTION_KEY_ENV, base64.b64encode(b"x" * 32).decode())
|
||||
with pytest.raises(RuntimeError, match="not a usable Fernet key"):
|
||||
load_key()
|
||||
|
||||
|
||||
def test_load_key_rejects_the_env_example_placeholder(monkeypatch):
|
||||
monkeypatch.setattr(settings, "ENVIRONMENT", "production")
|
||||
monkeypatch.setenv(ENCRYPTION_KEY_ENV, "your_fernet_encryption_key_here")
|
||||
with pytest.raises(RuntimeError, match="not a usable Fernet key"):
|
||||
load_key()
|
||||
|
||||
|
||||
def test_load_key_is_fatal_when_unset_outside_development(monkeypatch):
|
||||
monkeypatch.setattr(settings, "ENVIRONMENT", "production")
|
||||
monkeypatch.delenv(ENCRYPTION_KEY_ENV, raising=False)
|
||||
with pytest.raises(RuntimeError, match="is not set"):
|
||||
load_key()
|
||||
|
||||
|
||||
def test_load_key_falls_back_to_a_throwaway_key_in_development(monkeypatch):
|
||||
"""An unusable key must not stop local development — a generated one is fine
|
||||
there, since nothing durable is encrypted under it."""
|
||||
monkeypatch.setattr(settings, "ENVIRONMENT", "development")
|
||||
monkeypatch.setenv(ENCRYPTION_KEY_ENV, "your_fernet_encryption_key_here")
|
||||
Fernet(load_key()) # usable, and no raise
|
||||
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.core.error_handlers import (
|
||||
INTERNAL_ERROR_DETAIL,
|
||||
ClientSafeHTTPException,
|
||||
install_error_handlers,
|
||||
)
|
||||
|
||||
# What a leaked `detail=str(e)` looks like in practice: the text carries the
|
||||
# connection string, the SQL, the file path — none of it the caller's business.
|
||||
LEAKED = 'relation "users" does not exist at postgresql://admin:hunter2@10.0.0.4/app'
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
app = FastAPI()
|
||||
install_error_handlers(app)
|
||||
|
||||
@app.get("/boom")
|
||||
def boom():
|
||||
raise HTTPException(status_code=500, detail=LEAKED)
|
||||
|
||||
@app.get("/unavailable")
|
||||
def unavailable():
|
||||
raise HTTPException(status_code=503, detail=LEAKED)
|
||||
|
||||
@app.get("/missing")
|
||||
def missing():
|
||||
raise HTTPException(status_code=404, detail="Agent not found")
|
||||
|
||||
@app.get("/upstream")
|
||||
def upstream():
|
||||
raise ClientSafeHTTPException(status_code=502, detail="Recipient not on WhatsApp")
|
||||
|
||||
@app.get("/unauthorized")
|
||||
def unauthorized():
|
||||
raise HTTPException(
|
||||
status_code=401, detail="Not authenticated", headers={"WWW-Authenticate": "Bearer"}
|
||||
)
|
||||
|
||||
return TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
|
||||
class TestSanitizedHttpExceptions:
|
||||
def test_a_500_does_not_hand_the_caller_our_error_text(self, client):
|
||||
response = client.get("/boom")
|
||||
|
||||
assert response.status_code == 500
|
||||
assert response.json()["detail"] == INTERNAL_ERROR_DETAIL
|
||||
assert "postgresql://" not in response.text
|
||||
assert "users" not in response.text
|
||||
|
||||
def test_every_5xx_is_covered_not_just_500(self, client):
|
||||
response = client.get("/unavailable")
|
||||
|
||||
assert response.status_code == 503
|
||||
assert response.json()["detail"] == INTERNAL_ERROR_DETAIL
|
||||
|
||||
def test_the_real_reason_is_logged_with_the_route(self, client, caplog):
|
||||
with caplog.at_level("ERROR"):
|
||||
client.get("/boom")
|
||||
|
||||
assert LEAKED in caplog.text, "the original detail must survive in the log"
|
||||
assert "/boom" in caplog.text, "the log must say which route produced it"
|
||||
|
||||
def test_4xx_detail_is_written_for_the_caller_and_kept(self, client):
|
||||
response = client.get("/missing")
|
||||
|
||||
assert response.status_code == 404
|
||||
assert response.json()["detail"] == "Agent not found"
|
||||
|
||||
def test_4xx_headers_survive(self, client):
|
||||
response = client.get("/unauthorized")
|
||||
|
||||
assert response.status_code == 401
|
||||
assert response.headers["WWW-Authenticate"] == "Bearer"
|
||||
assert response.json()["detail"] == "Not authenticated"
|
||||
|
||||
def test_a_deliberate_5xx_message_still_reaches_the_caller(self, client):
|
||||
# Upstream wording the operator has to act on is useless in a log only.
|
||||
response = client.get("/upstream")
|
||||
|
||||
assert response.status_code == 502
|
||||
assert response.json()["detail"] == "Recipient not on WhatsApp"
|
||||
|
||||
def test_validation_errors_are_untouched(self, client):
|
||||
# 422 bodies name the offending field; they carry nothing internal.
|
||||
app = client.app
|
||||
|
||||
@app.get("/typed")
|
||||
def typed(count: int):
|
||||
return {"count": count}
|
||||
|
||||
response = TestClient(app).get("/typed", params={"count": "not-a-number"})
|
||||
|
||||
assert response.status_code == 422
|
||||
assert response.json()["detail"][0]["loc"] == ["query", "count"]
|
||||
@@ -0,0 +1,32 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
from app.core.external_urls import external_base
|
||||
|
||||
|
||||
def test_default_when_unset(monkeypatch):
|
||||
monkeypatch.delenv("TELEGRAM_API_BASE", raising=False)
|
||||
assert external_base("TELEGRAM_API_BASE", "https://api.telegram.org") == "https://api.telegram.org"
|
||||
|
||||
|
||||
def test_override_without_trailing_slash(monkeypatch):
|
||||
monkeypatch.setenv("TELEGRAM_API_BASE", "http://127.0.0.1:14020/telegram/")
|
||||
assert external_base("TELEGRAM_API_BASE", "https://api.telegram.org") == "http://127.0.0.1:14020/telegram"
|
||||
|
||||
|
||||
def test_blank_override_falls_back(monkeypatch):
|
||||
monkeypatch.setenv("TELEGRAM_API_BASE", "")
|
||||
assert external_base("TELEGRAM_API_BASE", "https://api.telegram.org") == "https://api.telegram.org"
|
||||
@@ -0,0 +1,100 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
"""
|
||||
|
||||
import io
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from starlette.datastructures import Headers, UploadFile
|
||||
|
||||
from app.core.file_validation import PDF_MAGIC, read_validated, safe_filename
|
||||
|
||||
PDF_TYPES = ("application/pdf",)
|
||||
|
||||
|
||||
def _pdf_upload(content: bytes, filename="doc.pdf", content_type="application/pdf"):
|
||||
headers = Headers({"content-type": content_type}) if content_type else None
|
||||
return UploadFile(io.BytesIO(content), filename=filename, headers=headers)
|
||||
|
||||
|
||||
# ---- safe_filename (path traversal) ----------------------------------------
|
||||
|
||||
@pytest.mark.parametrize("name", ["../../etc/passwd", "..\\..\\x.pdf", "/abs/path/a.pdf", "a/b/c.pdf"])
|
||||
def test_safe_filename_cannot_escape_target_dir(name):
|
||||
out = safe_filename(name)
|
||||
assert "/" not in out and "\\" not in out
|
||||
# The real property: joining under a base dir stays inside it.
|
||||
joined = os.path.normpath(os.path.join("temp", out))
|
||||
assert joined.startswith("temp" + os.sep)
|
||||
|
||||
|
||||
def test_safe_filename_prefixes_uuid_and_keeps_basename():
|
||||
out = safe_filename("report.pdf")
|
||||
assert out.endswith("_report.pdf")
|
||||
assert out != safe_filename("report.pdf") # unique per call
|
||||
|
||||
|
||||
def test_safe_filename_handles_empty_and_dots():
|
||||
assert safe_filename("").endswith("_upload")
|
||||
assert safe_filename("..").endswith("_upload")
|
||||
|
||||
|
||||
# ---- read_validated ---------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_validated_accepts_pdf():
|
||||
content = b"%PDF-1.7\n...body..."
|
||||
data = await read_validated(
|
||||
_pdf_upload(content), max_size=1024, allowed_content_types=PDF_TYPES,
|
||||
magic_prefix=PDF_MAGIC, label="PDF",
|
||||
)
|
||||
assert data == content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_validated_rejects_wrong_content_type():
|
||||
with pytest.raises(HTTPException) as e:
|
||||
await read_validated(
|
||||
_pdf_upload(b"%PDF-1.7", content_type="text/html"),
|
||||
max_size=1024, allowed_content_types=PDF_TYPES, magic_prefix=PDF_MAGIC, label="PDF",
|
||||
)
|
||||
assert e.value.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_validated_rejects_bad_magic_bytes():
|
||||
with pytest.raises(HTTPException) as e:
|
||||
await read_validated(
|
||||
_pdf_upload(b"<html>not a pdf</html>"),
|
||||
max_size=1024, allowed_content_types=PDF_TYPES, magic_prefix=PDF_MAGIC, label="PDF",
|
||||
)
|
||||
assert e.value.status_code == 400
|
||||
assert "not a valid" in e.value.detail.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_validated_rejects_oversize():
|
||||
with pytest.raises(HTTPException) as e:
|
||||
await read_validated(
|
||||
_pdf_upload(b"%PDF-" + b"x" * 5000),
|
||||
max_size=1024, allowed_content_types=PDF_TYPES, magic_prefix=PDF_MAGIC, label="PDF",
|
||||
)
|
||||
assert e.value.status_code == 413
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_validated_rejects_empty():
|
||||
with pytest.raises(HTTPException) as e:
|
||||
await read_validated(
|
||||
_pdf_upload(b""), max_size=1024, allowed_content_types=PDF_TYPES,
|
||||
magic_prefix=PDF_MAGIC, label="PDF",
|
||||
)
|
||||
assert e.value.status_code == 400
|
||||
@@ -0,0 +1,340 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
from app.knowledge.knowledge_base import KnowledgeManager
|
||||
from app.models.knowledge import Knowledge, SourceType
|
||||
from app.models.knowledge_to_agent import KnowledgeToAgent
|
||||
from uuid import UUID, uuid4
|
||||
import os
|
||||
from agno.vectordb.pgvector import PgVector, SearchType
|
||||
|
||||
@pytest.fixture
|
||||
def mock_db():
|
||||
"""Create a mock database session"""
|
||||
db = MagicMock()
|
||||
return db
|
||||
|
||||
@pytest.fixture
|
||||
def mock_ai_config():
|
||||
"""Create a mock AI config"""
|
||||
config = MagicMock()
|
||||
config.encrypted_api_key = "encrypted_key"
|
||||
return config
|
||||
|
||||
@pytest.fixture
|
||||
def mock_knowledge():
|
||||
"""Create a mock knowledge entry"""
|
||||
knowledge_id = uuid4()
|
||||
org_id = uuid4()
|
||||
return Knowledge(
|
||||
id=knowledge_id,
|
||||
organization_id=org_id,
|
||||
source="test_source",
|
||||
source_type=SourceType.FILE,
|
||||
table_name="test_table",
|
||||
schema="test_schema"
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_vector_db():
|
||||
"""Create a mock vector database"""
|
||||
vector_db = MagicMock(spec=PgVector)
|
||||
vector_db.table_name = "test_table"
|
||||
vector_db.schema = "test_schema"
|
||||
return vector_db
|
||||
|
||||
@pytest.fixture
|
||||
def knowledge_manager(mock_db, mock_ai_config, mock_vector_db):
|
||||
"""Create a KnowledgeManager instance with mocked dependencies"""
|
||||
org_id = uuid4()
|
||||
agent_id = uuid4()
|
||||
|
||||
with patch('app.knowledge.knowledge_base.SessionLocal') as mock_session_local, \
|
||||
patch('app.knowledge.knowledge_base.AIConfigRepository') as mock_ai_config_repo, \
|
||||
patch('app.knowledge.knowledge_base.OptimizedPgVector') as mock_pg_vector, \
|
||||
patch('app.knowledge.knowledge_base.KnowledgeRepository') as mock_knowledge_repo, \
|
||||
patch('app.knowledge.knowledge_base.KnowledgeToAgentRepository') as mock_link_repo:
|
||||
|
||||
# Configure basic mocks
|
||||
mock_session_local.return_value.__enter__.return_value = mock_db
|
||||
mock_ai_config_repo.return_value.get_active_config.return_value = mock_ai_config
|
||||
mock_pg_vector.return_value = mock_vector_db
|
||||
|
||||
# Configure repository mocks
|
||||
mock_knowledge_repo_instance = MagicMock()
|
||||
mock_link_repo_instance = MagicMock()
|
||||
mock_knowledge_repo.return_value = mock_knowledge_repo_instance
|
||||
mock_link_repo.return_value = mock_link_repo_instance
|
||||
|
||||
manager = KnowledgeManager(org_id=org_id, agent_id=str(agent_id))
|
||||
|
||||
# Replace repository instances with mocks
|
||||
manager.knowledge_repo = mock_knowledge_repo_instance
|
||||
manager.link_repo = mock_link_repo_instance
|
||||
|
||||
return manager
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_pdf_urls_success(knowledge_manager):
|
||||
"""A PDF URL is downloaded through the guarded fetcher, then indexed."""
|
||||
urls = ["https://example.com/docs/test.pdf"]
|
||||
|
||||
with patch('app.knowledge.knowledge_base.download_pdf') as mock_download, \
|
||||
patch.object(knowledge_manager, 'add_pdf_files', AsyncMock(return_value=True)) as mock_add_files, \
|
||||
patch('app.knowledge.knowledge_base.delete_file_from_s3', AsyncMock()) as mock_delete:
|
||||
result = await knowledge_manager.add_pdf_urls(urls)
|
||||
|
||||
assert result is True
|
||||
url, temp_path = mock_download.call_args.args
|
||||
assert url == urls[0]
|
||||
assert os.path.basename(temp_path) == "test.pdf"
|
||||
mock_add_files.assert_awaited_once_with([temp_path], filename="test")
|
||||
mock_delete.assert_not_called() # not one of our S3 uploads
|
||||
assert not os.path.exists(os.path.dirname(temp_path)) # temp dir cleaned up
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_pdf_urls_failure(knowledge_manager):
|
||||
"""A refused or failed download is reported, and nothing is indexed."""
|
||||
from app.knowledge.url_safety import BlockedHostError
|
||||
|
||||
with patch('app.knowledge.knowledge_base.download_pdf',
|
||||
side_effect=BlockedHostError("internal")), \
|
||||
patch.object(knowledge_manager, 'add_pdf_files', AsyncMock()) as mock_add_files:
|
||||
result = await knowledge_manager.add_pdf_urls(["http://169.254.169.254/latest"])
|
||||
|
||||
assert result is False
|
||||
mock_add_files.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize("url, name", [
|
||||
("https://example.com/a/guide.pdf", "guide"),
|
||||
("https://example.com/a/My%20Guide.pdf?x=1", "My Guide"),
|
||||
# An encoded traversal decodes to a separator; only the last part survives.
|
||||
("https://b.s3.amazonaws.com/knowledge/..%2F..%2Fapp%2Fmain.py", "main"),
|
||||
("https://example.com/", "example.com"),
|
||||
])
|
||||
def test_pdf_name_from_url_never_carries_a_path(url, name):
|
||||
from app.knowledge.knowledge_base import _pdf_name_from_url
|
||||
assert _pdf_name_from_url(url) == name
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_this_orgs_uploaded_pdf_is_deleted_from_s3(knowledge_manager):
|
||||
from app.knowledge.knowledge_base import settings
|
||||
own = f"https://{settings.S3_BUCKET}.s3.us-east-1.amazonaws.com/knowledge/{knowledge_manager.org_id}/f.pdf?X-Amz-Signature=x"
|
||||
other_org = f"https://{settings.S3_BUCKET}.s3.us-east-1.amazonaws.com/knowledge/{uuid4()}/f.pdf"
|
||||
# Another bucket, laid out so a path-style read lands on our org's key.
|
||||
foreign = f"https://evil.s3.amazonaws.com/x/knowledge/{knowledge_manager.org_id}/f.pdf"
|
||||
|
||||
with patch('app.knowledge.knowledge_base.delete_file_from_s3', AsyncMock(return_value=True)) as mock_delete:
|
||||
for url in (own, other_org, foreign):
|
||||
await knowledge_manager._delete_uploaded_pdf(url)
|
||||
|
||||
mock_delete.assert_awaited_once_with(own)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_websites_success(knowledge_manager, mock_knowledge):
|
||||
"""Test successful addition of websites"""
|
||||
# Setup
|
||||
urls = ["http://example.com"]
|
||||
max_links = 5
|
||||
|
||||
with patch('app.knowledge.knowledge_base.EnhancedWebsiteKnowledgeBase') as mock_web_kb, \
|
||||
patch.object(knowledge_manager, '_add_knowledge_source', return_value=mock_knowledge) as mock_add_source:
|
||||
# Configure mock
|
||||
mock_web_kb_instance = MagicMock()
|
||||
mock_web_kb.return_value = mock_web_kb_instance
|
||||
|
||||
# Execute
|
||||
result = await knowledge_manager.add_websites(urls, max_links)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
mock_web_kb.assert_called_once()
|
||||
mock_web_kb_instance.load.assert_called_once()
|
||||
mock_add_source.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_pdf_files_success(knowledge_manager, mock_knowledge):
|
||||
"""Test successful addition of PDF files"""
|
||||
# Setup
|
||||
files = ["/path/to/test.pdf"]
|
||||
|
||||
with patch('app.knowledge.knowledge_base.EnhancedPDFKnowledgeBase') as mock_pdf_kb, \
|
||||
patch('app.knowledge.knowledge_base.PDFImageReader') as mock_pdf_reader, \
|
||||
patch('os.path.exists') as mock_exists, \
|
||||
patch('os.path.getsize') as mock_getsize, \
|
||||
patch('os.path.basename') as mock_basename, \
|
||||
patch('os.path.splitext') as mock_splitext, \
|
||||
patch.object(knowledge_manager, '_add_knowledge_source', return_value=mock_knowledge) as mock_add_source:
|
||||
# Configure mocks
|
||||
mock_pdf_kb_instance = MagicMock()
|
||||
mock_pdf_kb.return_value = mock_pdf_kb_instance
|
||||
mock_exists.return_value = True
|
||||
mock_getsize.return_value = 1024
|
||||
mock_basename.return_value = "test.pdf"
|
||||
mock_splitext.return_value = ("test", ".pdf")
|
||||
|
||||
# Execute
|
||||
result = await knowledge_manager.add_pdf_files(files)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
mock_pdf_kb.assert_called_once()
|
||||
mock_pdf_kb_instance.load.assert_called_once()
|
||||
mock_add_source.assert_called_once()
|
||||
|
||||
def test_get_knowledge_base_for_agent(knowledge_manager, mock_knowledge):
|
||||
"""Test getting knowledge base for a specific agent"""
|
||||
# Setup
|
||||
mock_knowledge.agent_links = [
|
||||
MagicMock(agent_id=knowledge_manager.agent_id)
|
||||
]
|
||||
|
||||
with patch('app.knowledge.knowledge_base.KnowledgeRepository') as mock_knowledge_repo_class:
|
||||
mock_knowledge_repo_instance = MagicMock()
|
||||
mock_knowledge_repo_class.return_value = mock_knowledge_repo_instance
|
||||
mock_knowledge_repo_instance.get_by_agent.return_value = [mock_knowledge]
|
||||
|
||||
# Execute
|
||||
result = knowledge_manager.get_knowledge_base()
|
||||
|
||||
# Assert
|
||||
assert len(result) == 1
|
||||
assert result[0]['id'] == mock_knowledge.id
|
||||
assert result[0]['source'] == mock_knowledge.source
|
||||
assert result[0]['source_type'] == mock_knowledge.source_type
|
||||
mock_knowledge_repo_instance.get_by_agent.assert_called_once_with(knowledge_manager.agent_id)
|
||||
|
||||
def test_get_knowledge_base_for_org(knowledge_manager, mock_knowledge):
|
||||
"""Test getting knowledge base for an organization"""
|
||||
# Setup
|
||||
knowledge_manager.agent_id = None
|
||||
mock_knowledge.agent_links = []
|
||||
|
||||
with patch('app.knowledge.knowledge_base.KnowledgeRepository') as mock_knowledge_repo_class:
|
||||
mock_knowledge_repo_instance = MagicMock()
|
||||
mock_knowledge_repo_class.return_value = mock_knowledge_repo_instance
|
||||
mock_knowledge_repo_instance.get_by_org.return_value = [mock_knowledge]
|
||||
|
||||
# Execute
|
||||
result = knowledge_manager.get_knowledge_base()
|
||||
|
||||
# Assert
|
||||
assert len(result) == 1
|
||||
assert result[0]['id'] == mock_knowledge.id
|
||||
assert result[0]['source'] == mock_knowledge.source
|
||||
assert result[0]['source_type'] == mock_knowledge.source_type
|
||||
mock_knowledge_repo_instance.get_by_org.assert_called_once_with(knowledge_manager.org_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_knowledge_pdf_file(knowledge_manager):
|
||||
"""Test processing PDF file knowledge"""
|
||||
# Setup
|
||||
queue_item = MagicMock()
|
||||
queue_item.id = 123 # ensure a real integer ID to avoid DB adapter issues
|
||||
queue_item.source_type = "pdf_file"
|
||||
queue_item.source = "/path/to/test.pdf"
|
||||
|
||||
with patch.object(knowledge_manager, 'add_pdf_files') as mock_add_pdf_files, \
|
||||
patch('os.path.exists') as mock_exists, \
|
||||
patch('os.remove') as mock_remove, \
|
||||
patch('app.knowledge.knowledge_base.KnowledgeQueueRepository') as mock_queue_repo_cls:
|
||||
# Configure mocks
|
||||
mock_add_pdf_files.return_value = True
|
||||
mock_exists.return_value = True
|
||||
mock_queue_repo = MagicMock()
|
||||
mock_queue_repo.update_progress.return_value = True
|
||||
mock_queue_repo.update_status.return_value = True
|
||||
mock_queue_repo_cls.return_value = mock_queue_repo
|
||||
|
||||
# Execute
|
||||
result = await knowledge_manager.process_knowledge(queue_item)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
mock_add_pdf_files.assert_called_once_with([queue_item.source])
|
||||
mock_exists.assert_called_once_with(queue_item.source)
|
||||
mock_remove.assert_called_once_with(queue_item.source)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_knowledge_unsupported_type(knowledge_manager):
|
||||
"""Test processing unsupported knowledge type"""
|
||||
# Setup
|
||||
queue_item = MagicMock()
|
||||
queue_item.source_type = "unsupported_type"
|
||||
|
||||
# Execute and Assert
|
||||
with pytest.raises(ValueError, match="Unsupported source type: unsupported_type"):
|
||||
await knowledge_manager.process_knowledge(queue_item)
|
||||
|
||||
# ---- PDFs over a limit: refused with the reason, not "failed to process" ----
|
||||
|
||||
from app.knowledge.pdf_download import PdfTooLargeError # noqa: E402
|
||||
from app.knowledge.pdf_limits import PdfTooManyPages # noqa: E402
|
||||
from tests.knowledge.test_pdf_limits import pdf_with_pages # noqa: E402
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_pdf_over_the_page_limit_is_refused_before_it_is_read(knowledge_manager, tmp_path):
|
||||
path = tmp_path / "catalogue.pdf"
|
||||
path.write_bytes(pdf_with_pages(4))
|
||||
|
||||
with patch('app.core.config.settings.KB_PDF_MAX_PAGES', 3), \
|
||||
patch('app.knowledge.knowledge_base.EnhancedPDFKnowledgeBase') as reader, \
|
||||
pytest.raises(PdfTooManyPages, match="4 pages; the limit is 3"):
|
||||
await knowledge_manager.add_pdf_files([str(path)])
|
||||
|
||||
reader.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_pdf_link_over_the_size_limit_says_so_without_the_link(knowledge_manager):
|
||||
signed = "https://bucket.s3.amazonaws.com/knowledge/x/big.pdf?X-Amz-Signature=secret"
|
||||
with patch('app.knowledge.knowledge_base.download_pdf', side_effect=PdfTooLargeError(25 * 1024 * 1024)), \
|
||||
pytest.raises(PdfTooLargeError) as refused:
|
||||
await knowledge_manager.add_pdf_urls([signed])
|
||||
|
||||
assert "25 MB" in str(refused.value)
|
||||
assert "amazonaws" not in str(refused.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_refused_upload_is_reported_and_its_copy_removed(knowledge_manager, tmp_path):
|
||||
path = tmp_path / "0e9f_catalogue.pdf"
|
||||
path.write_bytes(pdf_with_pages(4))
|
||||
queue_item = MagicMock(id=123, source_type="pdf_file", source=str(path))
|
||||
|
||||
with patch('app.core.config.settings.KB_PDF_MAX_PAGES', 3), \
|
||||
patch('app.knowledge.knowledge_base.KnowledgeQueueRepository') as queue_repo, \
|
||||
pytest.raises(PdfTooManyPages):
|
||||
await knowledge_manager.process_knowledge(queue_item)
|
||||
|
||||
assert not path.exists()
|
||||
queue_repo.return_value.update_status.assert_not_called() # never reported as completed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("source_type,method", [("pdf_file", "add_pdf_files"), ("pdf_url", "add_pdf_urls")])
|
||||
async def test_a_pdf_that_fails_is_never_marked_completed_first(knowledge_manager, source_type, method):
|
||||
queue_item = MagicMock(id=123, source_type=source_type, source="https://example.com/x.pdf")
|
||||
|
||||
with patch.object(knowledge_manager, method, AsyncMock(return_value=False)), \
|
||||
patch('app.knowledge.knowledge_base.KnowledgeQueueRepository') as queue_repo, \
|
||||
pytest.raises(Exception, match="Failed to process knowledge source"):
|
||||
await knowledge_manager.process_knowledge(queue_item)
|
||||
|
||||
queue_repo.return_value.update_status.assert_not_called()
|
||||
@@ -0,0 +1,103 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
import app.core.oauth_state as oauth_state
|
||||
|
||||
|
||||
class FakeRedis:
|
||||
"""Just the calls the state store makes, with real get-then-delete semantics."""
|
||||
|
||||
def __init__(self):
|
||||
self.data = {}
|
||||
self.ttls = {}
|
||||
|
||||
def set(self, key, value, ex=None):
|
||||
self.data[key] = value
|
||||
self.ttls[key] = ex
|
||||
|
||||
def pipeline(self):
|
||||
return FakePipeline(self)
|
||||
|
||||
|
||||
class FakePipeline:
|
||||
def __init__(self, redis):
|
||||
self.redis = redis
|
||||
self.ops = []
|
||||
|
||||
def get(self, key):
|
||||
self.ops.append(lambda: self.redis.data.get(key))
|
||||
|
||||
def delete(self, key):
|
||||
self.ops.append(lambda: int(self.redis.data.pop(key, None) is not None))
|
||||
|
||||
def execute(self):
|
||||
return [op() for op in self.ops]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_redis(monkeypatch):
|
||||
monkeypatch.setattr(oauth_state, "get_redis", lambda: None)
|
||||
oauth_state._fallback.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis(monkeypatch):
|
||||
fake = FakeRedis()
|
||||
monkeypatch.setattr(oauth_state, "get_redis", lambda: fake)
|
||||
return fake
|
||||
|
||||
|
||||
def test_redis_state_is_shared_and_single_use(redis):
|
||||
oauth_state.store_state("jira", "s1", "org-1")
|
||||
assert redis.ttls["jira_oauth_state:s1"] == oauth_state.STATE_TTL_SECONDS
|
||||
assert oauth_state.pop_state("jira", "s1") == "org-1"
|
||||
assert oauth_state.pop_state("jira", "s1") is None
|
||||
|
||||
|
||||
def test_kinds_do_not_redeem_each_others_states(redis):
|
||||
oauth_state.store_state("slack", "s1", "org-1")
|
||||
assert oauth_state.pop_state("jira", "s1") is None
|
||||
assert oauth_state.pop_state("slack", "s1") == "org-1"
|
||||
|
||||
|
||||
def test_fallback_is_single_use_and_per_kind(no_redis):
|
||||
oauth_state.store_state("crm", "s1", "org-1:hubspot")
|
||||
assert oauth_state.pop_state("jira", "s1") is None
|
||||
assert oauth_state.pop_state("crm", "s1") == "org-1:hubspot"
|
||||
assert oauth_state.pop_state("crm", "s1") is None
|
||||
|
||||
|
||||
def test_fallback_state_expires(no_redis, monkeypatch):
|
||||
now = 1_000_000.0
|
||||
monkeypatch.setattr(oauth_state.time, "time", lambda: now)
|
||||
oauth_state.store_state("jira", "s1", "org-1")
|
||||
now += oauth_state.STATE_TTL_SECONDS + 1
|
||||
assert oauth_state.pop_state("jira", "s1") is None
|
||||
|
||||
|
||||
def test_fallback_prunes_expired_states(no_redis, monkeypatch):
|
||||
now = 1_000_000.0
|
||||
monkeypatch.setattr(oauth_state.time, "time", lambda: now)
|
||||
oauth_state.store_state("jira", "abandoned", "org-1")
|
||||
now += oauth_state.STATE_TTL_SECONDS + 1
|
||||
oauth_state.store_state("jira", "fresh", "org-1")
|
||||
assert ("jira", "abandoned") not in oauth_state._fallback
|
||||
|
||||
|
||||
def test_empty_state_is_never_valid(no_redis):
|
||||
assert oauth_state.pop_state("jira", "") is None
|
||||
@@ -0,0 +1,393 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
"""OUTBOUND_URL_POLICY=public_only at connect time: DNS rebinding and
|
||||
redirects into the internal network are refused on the address actually
|
||||
dialled; policy "any" installs nothing."""
|
||||
|
||||
import dataclasses
|
||||
import inspect
|
||||
import socket
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpcore
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from agno.models.openai import OpenAIChat
|
||||
from agno.tools.mcp import SSEClientParams, StreamableHTTPClientParams
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.streamable_http import streamablehttp_client
|
||||
|
||||
from app.core.config import settings
|
||||
from app.core.outbound_http import (
|
||||
AsyncGuardedBackend,
|
||||
GuardedBackend,
|
||||
guarded_async_http_client,
|
||||
guarded_http_client,
|
||||
)
|
||||
from app.core.outbound_policy import OutboundTargetBlocked, check_outbound_target
|
||||
|
||||
|
||||
def refused_by_policy(error: BaseException) -> bool:
|
||||
"""The guard surfaces as a connect error; the policy refusal is in its causes."""
|
||||
while error is not None:
|
||||
if isinstance(error, OutboundTargetBlocked):
|
||||
return True
|
||||
error = error.__cause__ or error.__context__
|
||||
return False
|
||||
from app.models.mcp_tool import MCPTransportType
|
||||
from app.services import db_connector_service
|
||||
from app.services.db_connector_service import DBConnectorConfig
|
||||
from app.tools.mcp_manager import MCPToolsManager
|
||||
from app.tools.mcp_outbound import guarded_mcp_http_client, remote_params_classes
|
||||
from app.utils.agno_utils import PinnedOpenAIChat, create_model
|
||||
|
||||
PUBLIC_IP = "93.184.216.34"
|
||||
OK = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 2\r\n\r\n{}"
|
||||
|
||||
|
||||
def _redirect(location: str) -> bytes:
|
||||
return f"HTTP/1.1 302 Found\r\nLocation: {location}\r\nContent-Length: 0\r\n\r\n".encode()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def public_only(monkeypatch):
|
||||
monkeypatch.setattr(settings, "OUTBOUND_URL_POLICY", "public_only")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def policy_any(monkeypatch):
|
||||
monkeypatch.setattr(settings, "OUTBOUND_URL_POLICY", "any")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dns(monkeypatch):
|
||||
"""Fake resolver: ``dns[name] = [answer, ...]`` answers successive lookups
|
||||
in turn (the last one repeats), so a rebinding record is two answers."""
|
||||
answers = {}
|
||||
real = socket.getaddrinfo
|
||||
|
||||
def getaddrinfo(host, port, *args, **kwargs):
|
||||
if host not in answers:
|
||||
return real(host, port, *args, **kwargs)
|
||||
queue = answers[host]
|
||||
address = queue.pop(0) if len(queue) > 1 else queue[0]
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", (address, port or 0))]
|
||||
|
||||
monkeypatch.setattr(socket, "getaddrinfo", getaddrinfo)
|
||||
return answers
|
||||
|
||||
|
||||
class _Stream:
|
||||
def __init__(self, backend, response):
|
||||
self._backend, self._buffer = backend, [response]
|
||||
|
||||
def _read(self):
|
||||
return self._buffer.pop(0) if self._buffer else b""
|
||||
|
||||
def _start_tls(self, server_hostname):
|
||||
self._backend.sni.append(server_hostname)
|
||||
return self
|
||||
|
||||
|
||||
class _SyncStream(_Stream, httpcore.NetworkStream):
|
||||
def read(self, max_bytes, timeout=None):
|
||||
return self._read()
|
||||
|
||||
def write(self, buffer, timeout=None):
|
||||
self._backend.sent += buffer
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
def start_tls(self, ssl_context, server_hostname=None, timeout=None):
|
||||
return self._start_tls(server_hostname)
|
||||
|
||||
|
||||
class _AsyncStream(_Stream, httpcore.AsyncNetworkStream):
|
||||
async def read(self, max_bytes, timeout=None):
|
||||
return self._read()
|
||||
|
||||
async def write(self, buffer, timeout=None):
|
||||
self._backend.sent += buffer
|
||||
|
||||
async def aclose(self):
|
||||
pass
|
||||
|
||||
async def start_tls(self, ssl_context, server_hostname=None, timeout=None):
|
||||
return self._start_tls(server_hostname)
|
||||
|
||||
|
||||
class _Recording:
|
||||
"""Stands in for the real socket layer under the guard: records what it
|
||||
was asked to dial and serves a canned response per dialled address."""
|
||||
|
||||
def __init__(self, responses):
|
||||
self.responses, self.dialled, self.sni, self.sent = responses, [], [], b""
|
||||
|
||||
|
||||
class RecordingBackend(_Recording, httpcore.NetworkBackend):
|
||||
def connect_tcp(self, host, port, **kwargs):
|
||||
self.dialled.append((host, port))
|
||||
return _SyncStream(self, self.responses[host])
|
||||
|
||||
|
||||
class AsyncRecordingBackend(_Recording, httpcore.AsyncNetworkBackend):
|
||||
async def connect_tcp(self, host, port, **kwargs):
|
||||
self.dialled.append((host, port))
|
||||
return _AsyncStream(self, self.responses[host])
|
||||
|
||||
|
||||
def _guard_of(client):
|
||||
return client._transport._pool._network_backend
|
||||
|
||||
|
||||
def _sync_client(responses, **kwargs):
|
||||
client = guarded_http_client(follow_redirects=True, **kwargs)
|
||||
recorder = RecordingBackend(responses)
|
||||
_guard_of(client)._inner = recorder
|
||||
return client, recorder
|
||||
|
||||
|
||||
def _async_client(responses):
|
||||
client = guarded_async_http_client(follow_redirects=True)
|
||||
recorder = AsyncRecordingBackend(responses)
|
||||
_guard_of(client)._inner = recorder
|
||||
return client, recorder
|
||||
|
||||
|
||||
class TestGuard:
|
||||
def test_public_target_dials_the_vetted_ip_with_sni_and_host_unchanged(self, public_only, dns):
|
||||
dns["api.example.com"] = [PUBLIC_IP]
|
||||
client, recorder = _sync_client({PUBLIC_IP: OK})
|
||||
|
||||
assert client.get("https://api.example.com/v1/models").status_code == 200
|
||||
assert recorder.dialled == [(PUBLIC_IP, 443)]
|
||||
assert recorder.sni == ["api.example.com"]
|
||||
assert b"\r\nHost: api.example.com\r\n" in recorder.sent
|
||||
|
||||
def test_rebinding_after_the_check_is_refused(self, public_only, dns):
|
||||
dns["rebind.example.com"] = [PUBLIC_IP, "127.0.0.1"]
|
||||
check_outbound_target("https://rebind.example.com/v1", "Base URL") # passes: public
|
||||
client, recorder = _sync_client({})
|
||||
|
||||
with pytest.raises(httpx.ConnectError) as exc:
|
||||
client.get("https://rebind.example.com/v1/models")
|
||||
assert refused_by_policy(exc.value)
|
||||
assert recorder.dialled == []
|
||||
|
||||
def test_any_internal_answer_refuses_the_host(self, public_only, monkeypatch):
|
||||
monkeypatch.setattr(socket, "getaddrinfo", lambda *a, **k: [
|
||||
(socket.AF_INET, socket.SOCK_STREAM, 6, "", (PUBLIC_IP, 443)),
|
||||
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.7", 443)),
|
||||
])
|
||||
client, recorder = _sync_client({})
|
||||
with pytest.raises(httpx.ConnectError) as exc:
|
||||
client.get("https://mixed.example.com/")
|
||||
assert refused_by_policy(exc.value)
|
||||
assert recorder.dialled == []
|
||||
|
||||
@pytest.mark.parametrize("location", [
|
||||
"http://127.0.0.1:8080/admin",
|
||||
"http://169.254.169.254/latest/meta-data/",
|
||||
"http://internal.example.com/",
|
||||
])
|
||||
def test_redirect_to_an_internal_address_is_refused(self, public_only, dns, location):
|
||||
dns["api.example.com"] = [PUBLIC_IP]
|
||||
dns["internal.example.com"] = ["10.0.0.5"]
|
||||
client, recorder = _sync_client({PUBLIC_IP: _redirect(location)})
|
||||
|
||||
with pytest.raises(httpx.ConnectError) as exc:
|
||||
client.get("https://api.example.com/v1/models")
|
||||
assert refused_by_policy(exc.value)
|
||||
assert recorder.dialled == [(PUBLIC_IP, 443)]
|
||||
|
||||
def test_redirect_to_another_public_host_is_followed(self, public_only, dns):
|
||||
dns["api.example.com"] = [PUBLIC_IP]
|
||||
dns["cdn.example.net"] = ["93.184.216.35"]
|
||||
client, recorder = _sync_client({
|
||||
PUBLIC_IP: _redirect("https://cdn.example.net/x"), "93.184.216.35": OK,
|
||||
})
|
||||
assert client.get("https://api.example.com/").status_code == 200
|
||||
assert recorder.dialled == [(PUBLIC_IP, 443), ("93.184.216.35", 443)]
|
||||
assert recorder.sni == ["api.example.com", "cdn.example.net"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_dials_the_vetted_ip(self, public_only, dns):
|
||||
dns["api.example.com"] = [PUBLIC_IP]
|
||||
client, recorder = _async_client({PUBLIC_IP: OK})
|
||||
|
||||
assert (await client.get("https://api.example.com/v1")).status_code == 200
|
||||
assert recorder.dialled == [(PUBLIC_IP, 443)]
|
||||
assert recorder.sni == ["api.example.com"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_redirect_to_metadata_is_refused(self, public_only, dns):
|
||||
dns["api.example.com"] = [PUBLIC_IP]
|
||||
client, recorder = _async_client({PUBLIC_IP: _redirect("http://169.254.169.254/")})
|
||||
|
||||
with pytest.raises(httpx.ConnectError) as exc:
|
||||
await client.get("https://api.example.com/v1")
|
||||
assert refused_by_policy(exc.value)
|
||||
assert recorder.dialled == [(PUBLIC_IP, 443)]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rebinding_is_refused(self, public_only, dns):
|
||||
dns["rebind.example.com"] = [PUBLIC_IP, "169.254.169.254"]
|
||||
check_outbound_target("https://rebind.example.com", "MCP server URL")
|
||||
client, recorder = _async_client({})
|
||||
|
||||
with pytest.raises(httpx.ConnectError) as exc:
|
||||
await client.get("https://rebind.example.com/mcp")
|
||||
assert refused_by_policy(exc.value)
|
||||
assert recorder.dialled == []
|
||||
|
||||
def test_unresolvable_host_is_a_connect_error(self, public_only, monkeypatch):
|
||||
def fail(*a, **k):
|
||||
raise socket.gaierror("nodename nor servname provided")
|
||||
monkeypatch.setattr(socket, "getaddrinfo", fail)
|
||||
client, _ = _sync_client({})
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
client.get("https://nowhere.example.com/")
|
||||
|
||||
|
||||
class TestPolicyAny:
|
||||
def test_no_guarded_clients(self, policy_any):
|
||||
assert guarded_http_client() is None
|
||||
assert guarded_async_http_client() is None
|
||||
|
||||
def test_ai_model_is_plain_openai_chat(self, policy_any):
|
||||
model = create_model("OPENAI_COMPATIBLE", "key", "m", base_url="http://127.0.0.1:11434/v1")
|
||||
assert type(model) is OpenAIChat
|
||||
|
||||
def test_mcp_params_are_agnos_own(self, policy_any):
|
||||
assert remote_params_classes() == (SSEClientParams, StreamableHTTPClientParams)
|
||||
|
||||
def test_database_is_dialled_by_name(self, policy_any):
|
||||
config = _db_config(host="db.internal")
|
||||
with patch.object(db_connector_service, "_driver_connect") as driver:
|
||||
with db_connector_service._connection(config):
|
||||
pass
|
||||
driver.assert_called_once_with(config, "db.internal", 5432, None)
|
||||
|
||||
|
||||
class TestAIModelWiring:
|
||||
def test_sync_and_async_clients_are_guarded(self, public_only, dns):
|
||||
dns["llm.example.com"] = [PUBLIC_IP]
|
||||
model = create_model("OPENAI_COMPATIBLE", "key", "m", base_url="https://llm.example.com/v1")
|
||||
|
||||
assert isinstance(model, PinnedOpenAIChat)
|
||||
assert isinstance(_guard_of(model.get_client()._client), GuardedBackend)
|
||||
assert isinstance(_guard_of(model.get_async_client()._client), AsyncGuardedBackend)
|
||||
|
||||
def test_sdk_redirect_to_metadata_is_refused(self, public_only, dns):
|
||||
dns["llm.example.com"] = [PUBLIC_IP]
|
||||
model = PinnedOpenAIChat(api_key="k", id="m", base_url="https://llm.example.com/v1", max_retries=0)
|
||||
client = model.get_client()
|
||||
recorder = RecordingBackend({PUBLIC_IP: _redirect("http://169.254.169.254/latest/")})
|
||||
_guard_of(client._client)._inner = recorder
|
||||
|
||||
with pytest.raises(openai.APIConnectionError) as exc:
|
||||
client.models.list()
|
||||
assert refused_by_policy(exc.value)
|
||||
assert recorder.dialled == [(PUBLIC_IP, 443)]
|
||||
|
||||
|
||||
class TestMCPWiring:
|
||||
def _config(self, transport):
|
||||
return SimpleNamespace(
|
||||
name="Remote", transport_type=transport, url="https://mcp.example.com/mcp",
|
||||
headers={"Authorization": "Bearer t"}, timeout=None, sse_read_timeout=None,
|
||||
terminate_on_close=True,
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("transport", [MCPTransportType.HTTP, MCPTransportType.SSE])
|
||||
def test_remote_tool_carries_the_guarded_factory(self, public_only, dns, transport):
|
||||
dns["mcp.example.com"] = [PUBLIC_IP]
|
||||
tool = MCPToolsManager._build_remote_tool(self._config(transport))
|
||||
|
||||
# agno hands these to the mcp SDK as keyword arguments.
|
||||
params = dataclasses.asdict(tool.server_params)
|
||||
assert params["httpx_client_factory"] is guarded_mcp_http_client
|
||||
|
||||
def test_factory_builds_a_guarded_redirect_following_client(self, public_only):
|
||||
client = guarded_mcp_http_client(headers={"X": "1"}, timeout=httpx.Timeout(5))
|
||||
assert isinstance(_guard_of(client), AsyncGuardedBackend)
|
||||
assert client.follow_redirects is True
|
||||
assert client.headers["X"] == "1"
|
||||
|
||||
@pytest.mark.parametrize("transport_client", [sse_client, streamablehttp_client])
|
||||
def test_installed_mcp_sdk_accepts_the_factory(self, transport_client):
|
||||
assert "httpx_client_factory" in inspect.signature(transport_client).parameters
|
||||
|
||||
|
||||
def _db_config(**overrides):
|
||||
values = dict(
|
||||
id=None, organization_id=None, name="c", engine="postgresql",
|
||||
host="db.example.com", port=5432, database="d", username="u", password="p",
|
||||
)
|
||||
values.update(overrides)
|
||||
return DBConnectorConfig(**values)
|
||||
|
||||
|
||||
class TestDatabasePinning:
|
||||
def test_postgres_dials_the_vetted_ip_and_keeps_the_name(self, public_only, dns, monkeypatch):
|
||||
dns["db.example.com"] = [PUBLIC_IP]
|
||||
connect = MagicMock()
|
||||
monkeypatch.setattr("psycopg.connect", connect)
|
||||
with db_connector_service._connection(_db_config()):
|
||||
pass
|
||||
kwargs = connect.call_args.kwargs
|
||||
assert kwargs["hostaddr"] == PUBLIC_IP
|
||||
assert kwargs["host"] == "db.example.com"
|
||||
|
||||
def test_postgres_rebinding_is_refused(self, public_only, dns, monkeypatch):
|
||||
dns["db.example.com"] = [PUBLIC_IP, "127.0.0.1"]
|
||||
check_outbound_target("db.example.com", "Database host") # the save-time check
|
||||
connect = MagicMock()
|
||||
monkeypatch.setattr("psycopg.connect", connect)
|
||||
with pytest.raises(OutboundTargetBlocked, match="Database host"):
|
||||
with db_connector_service._connection(_db_config()):
|
||||
pass
|
||||
connect.assert_not_called()
|
||||
|
||||
def test_mysql_dials_the_vetted_ip(self, public_only, dns, monkeypatch):
|
||||
dns["db.example.com"] = [PUBLIC_IP]
|
||||
pymysql = MagicMock()
|
||||
monkeypatch.setitem(sys.modules, "pymysql", pymysql)
|
||||
with db_connector_service._connection(_db_config(engine="mysql", port=3306)):
|
||||
pass
|
||||
assert pymysql.connect.call_args.kwargs["host"] == PUBLIC_IP
|
||||
|
||||
def test_ssh_bastion_is_dialled_by_vetted_ip(self, public_only, dns, monkeypatch):
|
||||
dns["bastion.example.com"] = [PUBLIC_IP]
|
||||
sshtunnel = MagicMock()
|
||||
sshtunnel.SSHTunnelForwarder.return_value.local_bind_port = 40000
|
||||
monkeypatch.setitem(sys.modules, "sshtunnel", sshtunnel)
|
||||
config = _db_config(
|
||||
host="10.0.0.5", ssh_enabled=True, ssh_host="bastion.example.com",
|
||||
ssh_username="u", ssh_password="pw",
|
||||
)
|
||||
with patch.object(db_connector_service, "_driver_connect") as driver:
|
||||
with db_connector_service._connection(config):
|
||||
pass
|
||||
assert sshtunnel.SSHTunnelForwarder.call_args.args[0] == (PUBLIC_IP, 22)
|
||||
# The database itself sits behind the bastion and is reached through it.
|
||||
assert sshtunnel.SSHTunnelForwarder.call_args.kwargs["remote_bind_address"] == ("10.0.0.5", 5432)
|
||||
driver.assert_called_once_with(config, "127.0.0.1", 40000)
|
||||
@@ -0,0 +1,159 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from app.core.config import settings
|
||||
from app.core.outbound_policy import (
|
||||
OutboundTargetBlocked,
|
||||
check_outbound_target,
|
||||
require_public_target,
|
||||
)
|
||||
from app.knowledge import url_safety
|
||||
from app.models.mcp_tool import MCPTransportType
|
||||
from app.services import db_connector_service
|
||||
from app.services.db_connector_service import DBConnectorConfig, first_hop
|
||||
from app.tools.mcp_manager import MCPToolsManager
|
||||
from app.utils.agno_utils import create_model
|
||||
|
||||
INTERNAL_TARGETS = [
|
||||
"http://127.0.0.1:11434/v1",
|
||||
"http://localhost:11434/v1",
|
||||
"localhost:11434/v1", # no scheme: must not parse as hostless
|
||||
"http://169.254.169.254/latest/meta-data/",
|
||||
"https://10.1.2.3/mcp",
|
||||
"192.168.0.10", # bare database host
|
||||
"::1", # bare IPv6 host
|
||||
"http://[fd00::1]/mcp",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def public_only(monkeypatch):
|
||||
monkeypatch.setattr(settings, "OUTBOUND_URL_POLICY", "public_only")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def policy_any(monkeypatch):
|
||||
monkeypatch.setattr(settings, "OUTBOUND_URL_POLICY", "any")
|
||||
|
||||
|
||||
def _resolves_to(address):
|
||||
return patch.object(url_safety.socket, "getaddrinfo",
|
||||
return_value=[(2, 1, 6, "", (address, 0))])
|
||||
|
||||
|
||||
class TestPolicy:
|
||||
def test_default_is_any(self):
|
||||
assert settings.OUTBOUND_URL_POLICY == "any"
|
||||
|
||||
@pytest.mark.parametrize("target", INTERNAL_TARGETS)
|
||||
def test_any_allows_internal_targets(self, policy_any, target):
|
||||
check_outbound_target(target, "Target")
|
||||
|
||||
@pytest.mark.parametrize("target", INTERNAL_TARGETS)
|
||||
def test_public_only_refuses_internal_targets(self, public_only, target):
|
||||
with pytest.raises(OutboundTargetBlocked, match="Target must be a public address"):
|
||||
check_outbound_target(target, "Target")
|
||||
|
||||
def test_public_only_refuses_a_name_resolving_inside(self, public_only):
|
||||
with _resolves_to("10.0.0.9"), pytest.raises(OutboundTargetBlocked):
|
||||
check_outbound_target("db.internal.example.com", "Database host")
|
||||
|
||||
def test_public_only_allows_a_public_target(self, public_only):
|
||||
with _resolves_to("93.184.216.34"):
|
||||
check_outbound_target("https://api.example.com/v1", "Base URL")
|
||||
check_outbound_target("db.example.com", "Database host")
|
||||
|
||||
def test_an_empty_target_is_left_to_field_validation(self, public_only):
|
||||
check_outbound_target(None, "Base URL")
|
||||
check_outbound_target("", "Base URL")
|
||||
|
||||
def test_require_public_target_is_a_400(self, public_only):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
require_public_target("http://127.0.0.1", "MCP server URL")
|
||||
assert exc.value.status_code == 400
|
||||
assert "MCP server URL must be a public address" in exc.value.detail
|
||||
|
||||
|
||||
class TestCheckedAgainBeforeConnecting:
|
||||
"""Saving is not enough: a hostname can be re-pointed after it was saved."""
|
||||
|
||||
def test_ai_base_url(self, public_only):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
create_model("OPENAI_COMPATIBLE", "key", "m", base_url="http://127.0.0.1:11434/v1")
|
||||
assert "must be a public address" in exc.value.detail
|
||||
|
||||
def test_ai_base_url_allowed_under_any(self, policy_any):
|
||||
model = create_model("OPENAI_COMPATIBLE", "key", "m", base_url="http://127.0.0.1:11434/v1")
|
||||
assert model.base_url == "http://127.0.0.1:11434/v1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_server_url(self, public_only):
|
||||
config = SimpleNamespace(
|
||||
name="Internal", transport_type=MCPTransportType.HTTP,
|
||||
url="http://169.254.169.254/mcp", headers=None, timeout=None,
|
||||
sse_read_timeout=None, terminate_on_close=True,
|
||||
)
|
||||
with patch("app.tools.mcp_manager.MCPTools") as mcp_tools:
|
||||
result = await MCPToolsManager(connect_budget=5).test_tool_config(config)
|
||||
|
||||
assert result["success"] is False
|
||||
assert "MCP server URL must be a public address" in result["error"]
|
||||
mcp_tools.assert_not_called()
|
||||
|
||||
def _db_config(self, **overrides):
|
||||
values = dict(
|
||||
id=None, organization_id=None, name="c", engine="postgresql",
|
||||
host="10.0.0.5", port=5432, database="d", username="u", password="p",
|
||||
)
|
||||
values.update(overrides)
|
||||
return DBConnectorConfig(**values)
|
||||
|
||||
def test_database_host(self, public_only):
|
||||
with patch.object(db_connector_service, "_driver_connect") as driver:
|
||||
with pytest.raises(OutboundTargetBlocked, match="Database host"):
|
||||
with db_connector_service._connection(self._db_config()):
|
||||
pass
|
||||
driver.assert_not_called()
|
||||
|
||||
def test_ssh_host_is_the_first_hop_when_tunnelling(self, public_only):
|
||||
config = self._db_config(ssh_enabled=True, ssh_host="127.0.0.1", ssh_username="u")
|
||||
with patch.object(db_connector_service, "_ssh_tunnel") as tunnel:
|
||||
with pytest.raises(OutboundTargetBlocked, match="SSH host"):
|
||||
with db_connector_service._connection(config):
|
||||
pass
|
||||
tunnel.assert_not_called()
|
||||
|
||||
def test_private_database_behind_a_public_bastion_is_allowed(self, public_only):
|
||||
config = self._db_config(ssh_enabled=True, ssh_host="bastion.example.com", ssh_username="u")
|
||||
tunnel = MagicMock()
|
||||
tunnel.return_value.__enter__.return_value = ("127.0.0.1", 40000)
|
||||
with _resolves_to("93.184.216.34"), \
|
||||
patch.object(db_connector_service, "_ssh_tunnel", tunnel), \
|
||||
patch.object(db_connector_service, "_driver_connect") as driver:
|
||||
with db_connector_service._connection(config):
|
||||
pass
|
||||
driver.assert_called_once_with(config, "127.0.0.1", 40000)
|
||||
|
||||
|
||||
def test_first_hop_picks_the_bastion_only_when_tunnelling():
|
||||
assert first_hop(False, "bastion", "db") == ("db", "Database host")
|
||||
assert first_hop(True, "bastion", "db") == ("bastion", "SSH host")
|
||||
@@ -0,0 +1,43 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.core.params import MAX_DB_INT, MAX_PAGE, DbIntId, PageNumber
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
|
||||
@app.get("/things/{thing_id}")
|
||||
def read_thing(thing_id: DbIntId, page: PageNumber = 1):
|
||||
return {"thing_id": thing_id, "page": page}
|
||||
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
|
||||
def test_ids_the_database_can_hold_pass_through():
|
||||
assert client.get(f"/things/{MAX_DB_INT}?page={MAX_PAGE}").json() == {"thing_id": MAX_DB_INT, "page": MAX_PAGE}
|
||||
|
||||
|
||||
def test_ids_and_pages_the_database_would_choke_on_are_422():
|
||||
"""Past Postgres INTEGER, or a page whose OFFSET overflows, the query itself
|
||||
would fail with a 500; the request is refused up front instead."""
|
||||
assert client.get(f"/things/{MAX_DB_INT + 1}").status_code == 422
|
||||
assert client.get("/things/0").status_code == 422
|
||||
assert client.get(f"/things/1?page={MAX_PAGE + 1}").status_code == 422
|
||||
assert client.get("/things/1?page=0").status_code == 422
|
||||
@@ -0,0 +1,157 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
@patch('app.core.redis.redis')
|
||||
@patch('app.core.redis.settings')
|
||||
def test_init_redis_disabled(mock_settings, mock_redis):
|
||||
"""Test Redis initialization when Redis is disabled"""
|
||||
from app.core.redis import init_redis
|
||||
|
||||
# Set up the mock
|
||||
mock_settings.REDIS_ENABLED = False
|
||||
|
||||
# Call the function
|
||||
result = init_redis()
|
||||
|
||||
# Verify that no Redis connection was created
|
||||
assert result is None
|
||||
mock_redis.from_url.assert_not_called()
|
||||
|
||||
@patch('app.core.redis.redis')
|
||||
@patch('app.core.redis.settings')
|
||||
def test_init_redis_standard_url(mock_settings, mock_redis):
|
||||
"""Test Redis initialization with a standard Redis URL"""
|
||||
from app.core.redis import init_redis
|
||||
|
||||
# Set up the mocks
|
||||
mock_settings.REDIS_ENABLED = True
|
||||
mock_settings.REDIS_URL = "redis://localhost:6379"
|
||||
mock_redis_client = MagicMock()
|
||||
mock_redis.from_url.return_value = mock_redis_client
|
||||
|
||||
# Call the function
|
||||
result = init_redis()
|
||||
|
||||
# Verify Redis was initialized with the correct parameters
|
||||
mock_redis.from_url.assert_called_once_with(
|
||||
"redis://localhost:6379",
|
||||
socket_timeout=5.0,
|
||||
socket_connect_timeout=5.0,
|
||||
decode_responses=True
|
||||
)
|
||||
mock_redis_client.ping.assert_called_once()
|
||||
assert result == mock_redis_client
|
||||
|
||||
@patch('app.core.redis.redis')
|
||||
@patch('app.core.redis.settings')
|
||||
def test_init_redis_aws_elasticache(mock_settings, mock_redis):
|
||||
"""Test Redis initialization with AWS ElastiCache URL (should use TLS)"""
|
||||
from app.core.redis import init_redis
|
||||
|
||||
# Set up the mocks
|
||||
mock_settings.REDIS_ENABLED = True
|
||||
mock_settings.REDIS_URL = "redis://my-cluster.cache.amazonaws.com:6379"
|
||||
mock_redis_client = MagicMock()
|
||||
mock_redis.from_url.return_value = mock_redis_client
|
||||
|
||||
# Call the function
|
||||
result = init_redis()
|
||||
|
||||
# Verify Redis was initialized with TLS URL
|
||||
mock_redis.from_url.assert_called_once_with(
|
||||
"rediss://my-cluster.cache.amazonaws.com:6379",
|
||||
socket_timeout=5.0,
|
||||
socket_connect_timeout=5.0,
|
||||
decode_responses=True
|
||||
)
|
||||
mock_redis_client.ping.assert_called_once()
|
||||
assert result == mock_redis_client
|
||||
|
||||
@patch('app.core.redis.redis')
|
||||
@patch('app.core.redis.settings')
|
||||
def test_init_redis_connection_error(mock_settings, mock_redis):
|
||||
"""Test Redis initialization with connection error"""
|
||||
from app.core.redis import init_redis
|
||||
|
||||
# Set up the mocks
|
||||
mock_settings.REDIS_ENABLED = True
|
||||
mock_settings.REDIS_URL = "redis://localhost:6379"
|
||||
mock_redis_client = MagicMock()
|
||||
mock_redis.from_url.return_value = mock_redis_client
|
||||
mock_redis_client.ping.side_effect = Exception("Connection error")
|
||||
|
||||
# Call the function
|
||||
result = init_redis()
|
||||
|
||||
# Verify Redis client was created but ping failed
|
||||
mock_redis.from_url.assert_called_once()
|
||||
mock_redis_client.ping.assert_called_once()
|
||||
assert result is None
|
||||
|
||||
@patch('app.core.redis.init_redis')
|
||||
@patch('app.core.redis.redis_client', None) # Reset the global redis_client
|
||||
@patch('app.core.redis.settings')
|
||||
def test_get_redis_first_call(mock_settings, mock_init_redis):
|
||||
"""Test get_redis when called for the first time"""
|
||||
from app.core.redis import get_redis
|
||||
|
||||
# Set up the mocks
|
||||
mock_settings.REDIS_ENABLED = True
|
||||
mock_redis_client = MagicMock()
|
||||
mock_init_redis.return_value = mock_redis_client
|
||||
|
||||
# Call the function
|
||||
result = get_redis()
|
||||
|
||||
# Verify init_redis was called and the result is correct
|
||||
mock_init_redis.assert_called_once()
|
||||
assert result == mock_redis_client
|
||||
|
||||
@patch('app.core.redis.redis_client', MagicMock()) # Simulate existing redis_client
|
||||
@patch('app.core.redis.init_redis')
|
||||
@patch('app.core.redis.settings')
|
||||
def test_get_redis_cached_client(mock_settings, mock_init_redis, monkeypatch):
|
||||
"""Test get_redis when a Redis client is already cached"""
|
||||
from app.core.redis import get_redis, redis_client
|
||||
|
||||
# Set up the mocks
|
||||
mock_settings.REDIS_ENABLED = True
|
||||
existing_client = redis_client # Use the mocked redis_client
|
||||
|
||||
# Call the function
|
||||
result = get_redis()
|
||||
|
||||
# Verify init_redis was not called and the existing client is returned
|
||||
mock_init_redis.assert_not_called()
|
||||
assert result == existing_client
|
||||
|
||||
@patch('app.core.redis.redis_client', None) # Reset the global redis_client
|
||||
@patch('app.core.redis.settings')
|
||||
def test_get_redis_disabled(mock_settings):
|
||||
"""Test get_redis when Redis is disabled"""
|
||||
from app.core.redis import get_redis
|
||||
|
||||
# Set up the mocks
|
||||
mock_settings.REDIS_ENABLED = False
|
||||
|
||||
# Call the function
|
||||
result = get_redis()
|
||||
|
||||
# Verify no Redis client was initialized or returned
|
||||
assert result is None
|
||||
@@ -0,0 +1,343 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
from fastapi import UploadFile, HTTPException
|
||||
from io import BytesIO
|
||||
from app.core.s3 import (
|
||||
S3_MAX_PRESIGN_SECONDS,
|
||||
delete_file_from_s3,
|
||||
get_s3_client,
|
||||
get_s3_signed_url,
|
||||
is_s3_url,
|
||||
strip_s3_signature,
|
||||
upload_file_to_s3,
|
||||
)
|
||||
from app.core.config import Settings, settings
|
||||
from botocore.exceptions import ClientError
|
||||
|
||||
|
||||
def test_get_s3_client():
|
||||
"""Test the get_s3_client function"""
|
||||
with patch('boto3.client') as mock_boto3_client:
|
||||
mock_client = MagicMock()
|
||||
mock_boto3_client.return_value = mock_client
|
||||
|
||||
# The client is lru_cached; drop any instance an earlier test built so
|
||||
# boto3.client is actually invoked here.
|
||||
get_s3_client.cache_clear()
|
||||
client = get_s3_client()
|
||||
|
||||
# Verify boto3.client was called with correct parameters
|
||||
mock_boto3_client.assert_called_once_with(
|
||||
's3',
|
||||
aws_access_key_id=settings.AWS_ACCESS_KEY_ID,
|
||||
aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY,
|
||||
region_name=settings.S3_REGION
|
||||
)
|
||||
|
||||
assert client == mock_client
|
||||
|
||||
|
||||
def test_get_s3_client_passes_none_when_credentials_unset():
|
||||
"""Unset credentials must reach boto3 as None, not ''.
|
||||
|
||||
boto3 only falls back to its default chain — and therefore to an EC2/ECS
|
||||
instance role — when both credential kwargs are None. An empty string is
|
||||
taken as an explicit (blank) credential, which signs requests that AWS
|
||||
rejects with InvalidAccessKeyId.
|
||||
"""
|
||||
with patch('boto3.client') as mock_boto3_client, \
|
||||
patch.object(settings, 'AWS_ACCESS_KEY_ID', None), \
|
||||
patch.object(settings, 'AWS_SECRET_ACCESS_KEY', None):
|
||||
get_s3_client.cache_clear()
|
||||
get_s3_client()
|
||||
|
||||
kwargs = mock_boto3_client.call_args.kwargs
|
||||
assert kwargs['aws_access_key_id'] is None
|
||||
assert kwargs['aws_secret_access_key'] is None
|
||||
|
||||
get_s3_client.cache_clear()
|
||||
|
||||
|
||||
def test_aws_credentials_default_to_none_not_empty_string():
|
||||
"""A blank/absent AWS_ACCESS_KEY_ID env var must land as None on settings."""
|
||||
with patch.dict(os.environ, {'AWS_ACCESS_KEY_ID': '', 'AWS_SECRET_ACCESS_KEY': ''}):
|
||||
fresh = Settings()
|
||||
assert fresh.AWS_ACCESS_KEY_ID is None
|
||||
assert fresh.AWS_SECRET_ACCESS_KEY is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_s3_signed_url_with_s3_storage_enabled():
|
||||
"""Test get_s3_signed_url when S3 storage is enabled"""
|
||||
test_s3_url = f"https://{settings.S3_BUCKET}.s3.{settings.S3_REGION}.amazonaws.com/test/file.jpg"
|
||||
expected_signed_url = "https://signed-url.example.com"
|
||||
|
||||
with patch('app.core.s3.settings.S3_FILE_STORAGE', True), \
|
||||
patch('app.core.s3.get_s3_client') as mock_get_client:
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.generate_presigned_url.return_value = expected_signed_url
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
result = await get_s3_signed_url(test_s3_url)
|
||||
|
||||
# Verify the client was called correctly
|
||||
mock_client.generate_presigned_url.assert_called_once_with(
|
||||
'get_object',
|
||||
Params={
|
||||
'Bucket': settings.S3_BUCKET,
|
||||
'Key': 'test/file.jpg'
|
||||
},
|
||||
ExpiresIn=settings.S3_PRESIGN_EXPIRY_SECONDS
|
||||
)
|
||||
|
||||
assert result == expected_signed_url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_s3_signed_url_clamps_expiry_to_s3_maximum():
|
||||
"""A caller asking for longer than S3 allows is clamped, not passed through.
|
||||
|
||||
AWS rejects SigV4 presigns above S3_MAX_PRESIGN_SECONDS with
|
||||
AuthorizationQueryParametersError, which 403s every signed image.
|
||||
"""
|
||||
test_s3_url = f"https://{settings.S3_BUCKET}.s3.{settings.S3_REGION}.amazonaws.com/test/file.jpg"
|
||||
|
||||
with patch('app.core.s3.settings.S3_FILE_STORAGE', True), \
|
||||
patch('app.core.s3.get_s3_client') as mock_get_client:
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
# The pre-fix default: 30 days.
|
||||
await get_s3_signed_url(test_s3_url, expiration=2592000)
|
||||
|
||||
assert mock_client.generate_presigned_url.call_args.kwargs['ExpiresIn'] == S3_MAX_PRESIGN_SECONDS
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stored,expected", [
|
||||
(
|
||||
"https://b.s3.us-east-1.amazonaws.com/k.png?AWSAccessKeyId=A&Signature=S&Expires=1",
|
||||
"https://b.s3.us-east-1.amazonaws.com/k.png",
|
||||
),
|
||||
(
|
||||
"https://b.s3.us-east-1.amazonaws.com/k.png",
|
||||
"https://b.s3.us-east-1.amazonaws.com/k.png",
|
||||
),
|
||||
("/api/v1/uploads/help_center/k.png", "/api/v1/uploads/help_center/k.png"),
|
||||
("data:image/png;base64,AAAA", "data:image/png;base64,AAAA"),
|
||||
(None, None),
|
||||
("", ""),
|
||||
])
|
||||
def test_strip_s3_signature(stored, expected):
|
||||
"""Signed URLs must never reach the database — clients round-trip response
|
||||
values straight back into update payloads."""
|
||||
assert strip_s3_signature(stored) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_s3_signed_url_with_s3_storage_disabled():
|
||||
"""Test get_s3_signed_url when S3 storage is disabled"""
|
||||
test_s3_url = "https://example.com/test/file.jpg"
|
||||
|
||||
with patch('app.core.s3.settings.S3_FILE_STORAGE', False):
|
||||
result = await get_s3_signed_url(test_s3_url)
|
||||
assert result == test_s3_url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_s3_signed_url_with_exception():
|
||||
"""Test get_s3_signed_url when an exception occurs"""
|
||||
test_s3_url = f"https://{settings.S3_BUCKET}.s3.{settings.S3_REGION}.amazonaws.com/test/file.jpg"
|
||||
|
||||
with patch('app.core.s3.settings.S3_FILE_STORAGE', True), \
|
||||
patch('app.core.s3.get_s3_client') as mock_get_client:
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.generate_presigned_url.side_effect = Exception("Test exception")
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
result = await get_s3_signed_url(test_s3_url)
|
||||
|
||||
# Should return the original URL on error
|
||||
assert result == test_s3_url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_to_s3_success():
|
||||
"""Test successful file upload to S3"""
|
||||
# File content as bytes (the function now expects bytes directly)
|
||||
file_content = b"test file content"
|
||||
|
||||
folder = "test-folder"
|
||||
filename = "test-file.txt"
|
||||
content_type = "text/plain"
|
||||
|
||||
expected_url = f"https://{settings.S3_BUCKET}.s3.{settings.S3_REGION}.amazonaws.com/{folder}/{filename}"
|
||||
|
||||
with patch('app.core.s3.get_s3_client') as mock_get_client:
|
||||
mock_client = MagicMock()
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
result = await upload_file_to_s3(file_content, folder, filename, content_type)
|
||||
|
||||
# Verify S3 client was called correctly
|
||||
mock_client.put_object.assert_called_once_with(
|
||||
Bucket=settings.S3_BUCKET,
|
||||
Key=f"{folder}/{filename}",
|
||||
Body=file_content,
|
||||
ContentType=content_type
|
||||
)
|
||||
|
||||
assert result == expected_url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_to_s3_without_content_type():
|
||||
"""Test file upload to S3 without specifying content type"""
|
||||
file_content = b"test file content"
|
||||
|
||||
folder = "test-folder"
|
||||
filename = "test-file.txt"
|
||||
|
||||
with patch('app.core.s3.get_s3_client') as mock_get_client:
|
||||
mock_client = MagicMock()
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
await upload_file_to_s3(file_content, folder, filename)
|
||||
|
||||
# Verify S3 client was called without ContentType
|
||||
mock_client.put_object.assert_called_once_with(
|
||||
Bucket=settings.S3_BUCKET,
|
||||
Key=f"{folder}/{filename}",
|
||||
Body=file_content
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error", [
|
||||
ClientError({'Error': {'Code': 'TestException', 'Message': 'Test error message'}}, 'PutObject'),
|
||||
Exception("Test exception"),
|
||||
])
|
||||
async def test_upload_file_to_s3_raises_instead_of_falling_back(error):
|
||||
"""A failed S3 upload must fail the request, not silently write to local disk.
|
||||
|
||||
The old fallback returned a /api/v1/uploads/... path the caller could not
|
||||
distinguish from success — durable-looking, but gone on the next container
|
||||
restart.
|
||||
"""
|
||||
with patch('app.core.s3.get_s3_client') as mock_get_client:
|
||||
mock_client = MagicMock()
|
||||
mock_client.put_object.side_effect = error
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await upload_file_to_s3(b"test content", "folder", "file.txt")
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_to_s3_does_not_create_the_bucket():
|
||||
"""Provisioning belongs in deployment; the runtime identity should not need
|
||||
s3:CreateBucket, nor pay a head_bucket round-trip on every upload."""
|
||||
with patch('app.core.s3.get_s3_client') as mock_get_client:
|
||||
mock_client = MagicMock()
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
await upload_file_to_s3(b"test content", "folder", "file.txt")
|
||||
|
||||
mock_client.head_bucket.assert_not_called()
|
||||
mock_client.create_bucket.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_file_from_s3_success():
|
||||
"""Test successful file deletion from S3"""
|
||||
test_s3_url = f"https://{settings.S3_BUCKET}.s3.{settings.S3_REGION}.amazonaws.com/test/file.jpg"
|
||||
|
||||
with patch('app.core.s3.get_s3_client') as mock_get_client:
|
||||
mock_client = MagicMock()
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
result = await delete_file_from_s3(test_s3_url)
|
||||
|
||||
# Verify S3 client was called correctly
|
||||
mock_client.delete_object.assert_called_once_with(
|
||||
Bucket=settings.S3_BUCKET,
|
||||
Key="test/file.jpg"
|
||||
)
|
||||
|
||||
assert result is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_file_from_s3_exception():
|
||||
"""Test file deletion from S3 with an exception"""
|
||||
test_s3_url = f"https://{settings.S3_BUCKET}.s3.{settings.S3_REGION}.amazonaws.com/test/file.jpg"
|
||||
|
||||
with patch('app.core.s3.get_s3_client') as mock_get_client:
|
||||
mock_client = MagicMock()
|
||||
mock_client.delete_object.side_effect = Exception("Test exception")
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
result = await delete_file_from_s3(test_s3_url)
|
||||
|
||||
assert result is False
|
||||
|
||||
# ---------- is_s3_url ----------
|
||||
|
||||
@pytest.mark.parametrize("url", [
|
||||
# Every shape AWS serves, including the regional virtual-hosted form
|
||||
# url_for_s3_key produces and a dotted bucket name.
|
||||
"https://s3.amazonaws.com/bucket/key.pdf",
|
||||
"https://s3.us-east-1.amazonaws.com/bucket/key.pdf",
|
||||
"https://s3-us-west-2.amazonaws.com/bucket/key.pdf",
|
||||
"https://bucket.s3.amazonaws.com/key.pdf",
|
||||
"https://bucket.s3.us-east-1.amazonaws.com/key.pdf",
|
||||
"https://my.bucket.s3.us-east-1.amazonaws.com/key.pdf",
|
||||
"https://BUCKET.S3.AMAZONAWS.COM/key.pdf",
|
||||
])
|
||||
def test_is_s3_url_accepts_every_s3_endpoint_shape(url):
|
||||
assert is_s3_url(url) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("url", [
|
||||
None, "",
|
||||
# The whole point: an attacker-supplied URL can carry the string anywhere
|
||||
# outside the host, and a substring check would call these ours.
|
||||
"https://evil.com/s3.amazonaws.com/key.pdf",
|
||||
"https://evil.com/?x=s3.amazonaws.com",
|
||||
"https://evil.com/#s3.amazonaws.com",
|
||||
"https://s3.amazonaws.com.attacker.com/key.pdf",
|
||||
"https://foos3.amazonaws.com/key.pdf",
|
||||
"https://notamazonaws.com/key.pdf",
|
||||
"https://example.com/file.pdf",
|
||||
])
|
||||
def test_is_s3_url_rejects_lookalikes(url):
|
||||
assert is_s3_url(url) is False
|
||||
|
||||
|
||||
def test_is_s3_url_matches_what_url_for_s3_key_produces():
|
||||
"""Round-trip guard: whatever we generate must be recognised as ours."""
|
||||
from app.core.s3 import url_for_s3_key
|
||||
|
||||
assert is_s3_url(url_for_s3_key("some/key.pdf")) is True
|
||||
@@ -0,0 +1,127 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock
|
||||
import socketio
|
||||
|
||||
from app.core.socketio import configure_socketio, sio, socket_app
|
||||
|
||||
|
||||
def test_configure_socketio_with_cors():
|
||||
"""Test configuring socketio with CORS origins"""
|
||||
# Test with a string
|
||||
with patch.object(sio.eio, 'cors_allowed_origins', None):
|
||||
configure_socketio(cors_origins="https://example.com")
|
||||
assert sio.eio.cors_allowed_origins == ["https://example.com"]
|
||||
|
||||
# Test with a list
|
||||
with patch.object(sio.eio, 'cors_allowed_origins', None):
|
||||
configure_socketio(cors_origins=["https://example1.com", "https://example2.com"])
|
||||
assert sio.eio.cors_allowed_origins == ["https://example1.com", "https://example2.com"]
|
||||
|
||||
# Test with a set
|
||||
with patch.object(sio.eio, 'cors_allowed_origins', None):
|
||||
configure_socketio(cors_origins={"https://example1.com", "https://example2.com"})
|
||||
assert isinstance(sio.eio.cors_allowed_origins, list)
|
||||
assert set(sio.eio.cors_allowed_origins) == {"https://example1.com", "https://example2.com"}
|
||||
|
||||
|
||||
REDIS_OPTIONS = {
|
||||
'retry_on_timeout': True,
|
||||
'health_check_interval': 30,
|
||||
'socket_timeout': 5.0,
|
||||
'socket_connect_timeout': 5.0,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def restore_manager():
|
||||
"""configure_socketio swaps the shared server's manager; put the real one back."""
|
||||
manager, initialized = sio.manager, sio.manager_initialized
|
||||
yield
|
||||
sio.manager, sio.manager_initialized = manager, initialized
|
||||
|
||||
|
||||
@patch('app.core.socketio.settings')
|
||||
def test_configure_socketio_with_redis_disabled(mock_settings, restore_manager):
|
||||
"""Without Redis the server keeps its in-process manager."""
|
||||
mock_settings.REDIS_ENABLED = False
|
||||
manager = sio.manager
|
||||
|
||||
configure_socketio()
|
||||
|
||||
assert sio.manager is manager
|
||||
assert not isinstance(sio.manager, socketio.AsyncRedisManager)
|
||||
|
||||
|
||||
@patch('app.core.socketio.settings')
|
||||
@patch('socketio.AsyncRedisManager')
|
||||
def test_configure_socketio_with_redis_enabled(mock_redis_manager, mock_settings, restore_manager):
|
||||
"""With Redis, the Redis manager is the one the server emits through —
|
||||
not merely assigned to an attribute it never reads."""
|
||||
mock_settings.REDIS_ENABLED = True
|
||||
mock_settings.REDIS_URL = "redis://localhost:6379"
|
||||
mock_redis_instance = MagicMock()
|
||||
mock_redis_manager.return_value = mock_redis_instance
|
||||
|
||||
configure_socketio()
|
||||
|
||||
mock_redis_manager.assert_called_once_with(
|
||||
"redis://localhost:6379", write_only=False, channel='replypier', redis_options=REDIS_OPTIONS,
|
||||
)
|
||||
assert sio.manager is mock_redis_instance
|
||||
mock_redis_instance.set_server.assert_called_once_with(sio)
|
||||
assert sio.manager_initialized is False
|
||||
|
||||
|
||||
@patch('app.core.socketio.settings')
|
||||
@patch('socketio.AsyncRedisManager')
|
||||
def test_configure_socketio_with_elasticache(mock_redis_manager, mock_settings, restore_manager):
|
||||
"""ElastiCache URLs are upgraded to TLS."""
|
||||
mock_settings.REDIS_ENABLED = True
|
||||
mock_settings.REDIS_URL = "redis://my-cluster.cache.amazonaws.com:6379"
|
||||
mock_redis_manager.return_value = MagicMock()
|
||||
|
||||
configure_socketio()
|
||||
|
||||
mock_redis_manager.assert_called_once_with(
|
||||
"rediss://my-cluster.cache.amazonaws.com:6379", write_only=False, channel='replypier',
|
||||
redis_options={**REDIS_OPTIONS, 'ssl_cert_reqs': None},
|
||||
)
|
||||
|
||||
|
||||
@patch('app.core.socketio.settings')
|
||||
@patch('socketio.AsyncRedisManager')
|
||||
def test_configure_socketio_with_redis_exception(mock_redis_manager, mock_settings, restore_manager):
|
||||
"""A Redis manager that cannot be built leaves the in-process one serving."""
|
||||
mock_settings.REDIS_ENABLED = True
|
||||
mock_settings.REDIS_URL = "redis://localhost:6379"
|
||||
mock_redis_manager.side_effect = Exception("Redis connection error")
|
||||
manager = sio.manager
|
||||
|
||||
configure_socketio()
|
||||
|
||||
mock_redis_manager.assert_called_once()
|
||||
assert sio.manager is manager
|
||||
|
||||
|
||||
def test_plain_redis_gets_no_tls_options():
|
||||
"""redis-py rejects ssl_cert_reqs on redis://; with it the manager's listener
|
||||
fails in a loop and floods the log."""
|
||||
from app.core.socketio import redis_options
|
||||
assert 'ssl_cert_reqs' not in redis_options("redis://redis:6379/0")
|
||||
assert redis_options("rediss://cache.amazonaws.com:6379")['ssl_cert_reqs'] is None
|
||||
@@ -0,0 +1,59 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from app.core.socketio import needs_elasticache_tls
|
||||
|
||||
|
||||
class TestNeedsElasticacheTls:
|
||||
"""Whether a plaintext Redis URL points at ElastiCache, so it must be upgraded.
|
||||
|
||||
The old check asked whether ".cache.amazonaws.com" appeared anywhere in the
|
||||
URL, so a host that merely mentioned it in a query string or password passed.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize("url", [
|
||||
"redis://my-cluster.abc123.ng.0001.euw1.cache.amazonaws.com:6379/0",
|
||||
"redis://cache.amazonaws.com.cache.amazonaws.com:6379",
|
||||
])
|
||||
def test_real_elasticache_hosts_are_upgraded(self, url):
|
||||
assert needs_elasticache_tls(url) is True
|
||||
|
||||
@pytest.mark.parametrize("url", [
|
||||
# The host is the attacker's; the suffix only appears in the query.
|
||||
"redis://evil.example/?x=.cache.amazonaws.com",
|
||||
# ...or in the password.
|
||||
"redis://user:.cache.amazonaws.com@evil.example:6379/0",
|
||||
# ...or as a path segment.
|
||||
"redis://evil.example:6379/.cache.amazonaws.com",
|
||||
# A lookalike host that merely ends in a similar string.
|
||||
"redis://notcache.amazonaws.com.evil.example:6379",
|
||||
])
|
||||
def test_a_suffix_elsewhere_in_the_url_does_not_count(self, url):
|
||||
assert needs_elasticache_tls(url) is False
|
||||
|
||||
def test_a_plain_local_redis_is_left_alone(self):
|
||||
assert needs_elasticache_tls("redis://localhost:6379/0") is False
|
||||
|
||||
def test_an_already_tls_url_is_not_upgraded_again(self):
|
||||
assert needs_elasticache_tls(
|
||||
"rediss://my-cluster.euw1.cache.amazonaws.com:6379"
|
||||
) is False
|
||||
|
||||
@pytest.mark.parametrize("url", [None, "", "not a url"])
|
||||
def test_missing_or_junk_urls_are_handled(self, url):
|
||||
assert needs_elasticache_tls(url) is False
|
||||
@@ -0,0 +1,75 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
|
||||
Codebase-wide guard for a SQL construct that fails only at runtime.
|
||||
|
||||
`:param::type` looks like a bound parameter next to a Postgres cast, but
|
||||
SQLAlchemy's text() bind regex ends with a negative lookahead on ':', so the
|
||||
parameter is never recognised. The literal ':param' reaches the driver and
|
||||
Postgres rejects the statement with `syntax error at or near ":"`.
|
||||
|
||||
Nothing catches this before production: it compiles, it type-checks, and the
|
||||
SQLite test database never sees these statements. It broke knowledge link/unlink
|
||||
in production for five days. Use CAST(:param AS type) instead.
|
||||
"""
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
APP_DIR = Path(__file__).resolve().parents[2] / "app"
|
||||
|
||||
# ':name' immediately followed by '::' — the exact construct that silently
|
||||
# stops being a bind parameter.
|
||||
_PARAM_THEN_CAST = re.compile(r"(?<![:\w$]):([\w$]+)::")
|
||||
|
||||
|
||||
def _python_sources():
|
||||
return sorted(p for p in APP_DIR.rglob("*.py") if "__pycache__" not in p.parts)
|
||||
|
||||
|
||||
def test_no_bind_param_immediately_followed_by_a_cast():
|
||||
offenders = []
|
||||
for path in _python_sources():
|
||||
for lineno, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
|
||||
stripped = line.lstrip()
|
||||
# Skip prose: the fix is documented in comments that name the construct.
|
||||
if stripped.startswith("#"):
|
||||
continue
|
||||
for match in _PARAM_THEN_CAST.finditer(line):
|
||||
offenders.append(
|
||||
f"{path.relative_to(APP_DIR.parent)}:{lineno} -> :{match.group(1)}:: "
|
||||
f"(use CAST(:{match.group(1)} AS <type>))"
|
||||
)
|
||||
|
||||
assert not offenders, "SQL bind parameter swallowed by a '::' cast:\n" + "\n".join(offenders)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sql,expected",
|
||||
[
|
||||
("SELECT :a::jsonb", set()), # swallowed — the bug
|
||||
("SELECT CAST(:a AS jsonb)", {"a"}), # bound — the fix
|
||||
],
|
||||
)
|
||||
def test_the_guard_describes_real_sqlalchemy_behaviour(sql, expected):
|
||||
"""Pin the SQLAlchemy behaviour the guard above exists for, so the guard is
|
||||
removed deliberately if a future version ever starts parsing ':a::jsonb'.
|
||||
"""
|
||||
compiled = text(sql).compile(dialect=postgresql.dialect())
|
||||
assert set(compiled.params) == expected
|
||||
@@ -0,0 +1,43 @@
|
||||
"""
|
||||
Copyright 2024-2026 ReplyPier
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from app.channels.constants import WIDGET_CHANNELS, is_widget_channel
|
||||
|
||||
|
||||
class TestIsWidgetChannel:
|
||||
@pytest.mark.parametrize("channel", ["web", "shopify"])
|
||||
def test_our_own_widget_surfaces(self, channel):
|
||||
assert is_widget_channel(channel) is True
|
||||
|
||||
@pytest.mark.parametrize("channel", [None, ""])
|
||||
def test_absent_channel_means_the_widget(self, channel):
|
||||
"""Sessions predating the column, and anything that omits it."""
|
||||
assert is_widget_channel(channel) is True
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"channel",
|
||||
["telegram", "whatsapp", "messenger", "instagram", "slack", "email", "sms", "line"],
|
||||
)
|
||||
def test_external_platforms_are_not_the_widget(self, channel):
|
||||
assert is_widget_channel(channel) is False
|
||||
|
||||
def test_shopify_is_a_widget_channel(self):
|
||||
"""The whole point: labelling Shopify separately must not cost it the
|
||||
widget features that 'web' gets — the rating prompt above all."""
|
||||
assert "shopify" in WIDGET_CHANNELS
|
||||
assert is_widget_channel("shopify") is True
|
||||
Reference in new issue
Block a user