import asyncio from datetime import datetime import unittest from unittest.mock import AsyncMock, patch import httpx from fastapi import APIRouter, FastAPI, Request from client_source import current_client_source, resolve_client_source from upload_routes import FaceAnalysisFormRoute from test_ranking_consistency import database class SourceIsolationTest(unittest.IsolatedAsyncioTestCase): def setUp(self): database._RANK_CACHE.clear() database._LEADERBOARD_CACHE.clear() async def test_rank_sql_and_cache_are_partitioned_for_both_scopes(self): fetch = AsyncMock(side_effect=[{'t': 20, 'r': 4}, {'t': 10, 'r': 2}, {'t': 3, 'r': 1}, {'t': 2, 'r': 1}]) with patch.object(database, 'fetch_one', fetch): android = await database.get_face_score_rank(88, 'CN', client_source='android') wechat = await database.get_face_score_rank(88, 'CN', client_source='wechat') cached = await database.get_face_score_rank(88, 'CN', client_source='android') self.assertEqual(20, android['global_total']) self.assertEqual(3, wechat['global_total']) self.assertEqual(android, cached) self.assertEqual(4, fetch.await_count) for index, call in enumerate(fetch.await_args_list): self.assertIn('client_source = %s', call.args[0]) self.assertEqual('android' if index < 2 else 'wechat', call.args[1][-1]) self.assertNotIn('nickname LIKE', call.args[0]) async def test_leaderboard_cache_and_device_enrichment_do_not_cross_sources(self): async def fetch(sql, params): if 'tpl_app_user_devices' in sql: return [{'device_id': 'same', 'display_name': 'Android name'}] return [{'id': 1 if params[-1] == 'android' else 2, 'score': 88, 'nickname': 'android_release_same', 'region': 'CN'}] query = AsyncMock(side_effect=fetch) with patch.object(database, 'fetch_all', query): a = await database.get_face_score_leaderboard('CN', 50, 'same', 'android') w = await database.get_face_score_leaderboard('CN', 50, 'same', 'wechat') w2 = await database.get_face_score_leaderboard('CN', 50, 'same', 'wechat') self.assertEqual(1, a['items'][0]['record_id']) self.assertEqual(2, w['items'][0]['record_id']) self.assertEqual(w, w2) self.assertTrue(a['items'][0]['is_current_user']) self.assertFalse(w['items'][0]['is_current_user']) self.assertNotEqual('Android name', w['items'][0].get('display_name')) self.assertEqual(3, query.await_count) self.assertIn('candidate.client_source = %s', query.await_args_list[0].args[0]) self.assertIn('candidate.client_source = %s', query.await_args_list[2].args[0]) async def test_leaderboard_file_path_survives_serialization_and_cache(self): from models import LeaderboardResponse for source in ('android', 'wechat'): for region in (None, 'CN'): with self.subTest(source=source, region=region): path = f'face/{source}/portrait.webp' query = AsyncMock(return_value=[{'id': 7, 'score': 88, 'region': 'CN', 'file_path': path}]) with patch.object(database, 'fetch_all', query): fresh = await database.get_face_score_leaderboard(region, 50, client_source=source) cached = await database.get_face_score_leaderboard(region, 50, client_source=source) self.assertEqual(1, query.await_count) self.assertIn('candidate.file_path', query.await_args.args[0]) for result in (fresh, cached): payload = LeaderboardResponse(**result).model_dump() self.assertEqual(path, payload['items'][0]['file_path']) async def test_image_and_device_writes_include_source(self): execute = AsyncMock() with patch.object(database, 'execute', execute): for source in ('android', 'wechat'): token = current_client_source.set(source) try: await database.upsert_image_record(file_path='face.webp', category='face', nickname='same', score=88, is_cropped_face=True, size_bytes=1, last_modified=datetime.now(), bos_uploaded=True) finally: current_client_source.reset(token) self.assertEqual(source, execute.await_args.args[1][-1]) self.assertIn("IF(client_source = 'unknown'", execute.await_args.args[0]) await database.upsert_device_record(device_id='device') self.assertEqual('android', execute.await_args.args[1][-1]) async def test_concurrent_requests_and_background_tasks_keep_provenance(self): app = FastAPI() router = APIRouter(prefix='/facescore', route_class=FaceAnalysisFormRoute) tasks = [] seen = [] @router.get('/app/rank') @router.get('/wechat/rank') async def endpoint(request: Request): async def record(): await asyncio.sleep(.01) seen.append(current_client_source.get()) tasks.append(asyncio.create_task(record())) await asyncio.sleep(0) return {'source': current_client_source.get()} app.include_router(router) async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client: results = await asyncio.gather(client.get('/facescore/app/rank'), client.get('/facescore/wechat/rank')) await asyncio.gather(*tasks) self.assertEqual(['android', 'wechat'], [r.json()['source'] for r in results]) self.assertCountEqual(['android', 'wechat'], seen) self.assertEqual('unknown', current_client_source.get()) async def test_admin_filters_use_source_instead_of_region_or_nickname(self): fetch_all = AsyncMock(return_value=[]) fetch_one = AsyncMock(return_value={'total': 0}) with patch.object(database, 'fetch_all', fetch_all), patch.object(database, 'fetch_one', fetch_one): await database.fetch_paged_image_records(category=None, nickname=None, platform='miniprogram', offset=0, limit=20) await database.count_image_records(category=None, nickname=None, platform='miniprogram') await database.fetch_image_region_stats(category=None, nickname=None, platform='miniprogram') for call in [*fetch_all.await_args_list, *fetch_one.await_args_list]: self.assertIn('client_source = %s', call.args[0]) self.assertNotIn('region = %s', call.args[0]) self.assertEqual('wechat', call.args[1][0]) async def test_request_exception_restores_source_context(self): app = FastAPI() router = APIRouter(prefix='/facescore', route_class=FaceAnalysisFormRoute) @router.get('/app/failure') async def failure(): self.assertEqual('android', current_client_source.get()) raise RuntimeError('expected failure') app.include_router(router) async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client: with self.assertRaisesRegex(RuntimeError, 'expected failure'): await client.get('/facescore/app/failure') self.assertEqual('unknown', current_client_source.get()) def test_routes_override_client_headers_and_queries(self): for path, headers, expected in [ ('/facescore/app/analyze', {}, 'android'), ('/facescore/analyze', {'x-device-id': 'spoof'}, 'wechat'), ('/facescore/rank', {'x-device-id': 'legacy'}, 'android'), ('/facescore/rank', {}, 'wechat'), ('/facescore/wechat/leaderboard', {'x-device-id': 'spoof'}, 'wechat'), ('/facescore/web/analyze', {}, 'web'), ('/facescore/ios/analyze', {}, 'ios')]: request = Request({'type': 'http', 'path': path, 'query_string': b'client_source=android', 'headers': [(k.encode(), v.encode()) for k, v in headers.items()]}) self.assertEqual(expected, resolve_client_source(request)) if __name__ == '__main__': unittest.main()