fixed # 统一头像获取接口

tests/test_astrbot_avatars.py

@@ -0,0 +1,144 @@
+import asyncio
+import importlib.util
+from pathlib import Path
+import unittest
+from unittest.mock import AsyncMock, patch
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+class AvatarTests(unittest.IsolatedAsyncioTestCase):
+    def setUp(self):
+        spec = importlib.util.spec_from_file_location(
+            'farm_astrbot_avatars',
+            ROOT / 'integrations/astrbot/avatars.py',
+        )
+        self.module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(self.module)
+        self.cache = self.module.AvatarCache()
+
+    async def asyncTearDown(self):
+        await self.cache.close()
+
+    async def test_avatar_success_cache_and_concurrent_requests(self):
+        async def download(user_id):
+            await asyncio.sleep(0)
+            return b'avatar'
+
+        with patch.object(self.module, 'download_avatar',
+                          side_effect=download) as fetch:
+            results = await asyncio.gather(*(
+                self.cache.get('aiocqhttp', '123456') for _ in range(8)
+            ))
+            self.assertEqual(results, [b'avatar'] * 8)
+            self.assertEqual(
+                await self.cache.get('aiocqhttp', '123456'), b'avatar',
+            )
+            fetch.assert_awaited_once_with('123456')
+            self.assertEqual(
+                await self.cache.get('aiocqhttp', '654321'), b'avatar',
+            )
+            self.assertEqual(fetch.await_count, 2)
+
+    async def test_failure_negative_cache_and_expiry(self):
+        with patch.object(self.module, 'download_avatar',
+                          new=AsyncMock(side_effect=OSError('offline'))) as f:
+            with patch.object(self.module, 'monotonic', return_value=0):
+                self.assertIsNone(await self.cache.get('aiocqhttp', '12345'))
+            with patch.object(self.module, 'monotonic', return_value=20):
+                self.assertIsNone(await self.cache.get('aiocqhttp', '12345'))
+            self.assertEqual(f.await_count, 1)
+            with patch.object(self.module, 'monotonic', return_value=31):
+                self.assertIsNone(await self.cache.get('aiocqhttp', '12345'))
+            self.assertEqual(f.await_count, 2)
+
+    async def test_unknown_platform_invalid_id_and_closed_cache(self):
+        with patch.object(self.module, 'download_avatar') as fetch:
+            for platform, user in (
+                ('telegram', '12345'), ('qq_official', '12345'),
+                ('aiocqhttp', '../secret'), ('aiocqhttp', '１２３４５'),
+                ('aiocqhttp', ''), ('aiocqhttp', '1' * 21),
+            ):
+                self.assertIsNone(await self.cache.get(platform, user))
+            await self.cache.close()
+            self.assertIsNone(await self.cache.get('aiocqhttp', '12345'))
+            fetch.assert_not_called()
+
+    async def test_timeout_size_and_cache_bound(self):
+        async def slow(user_id):
+            await asyncio.sleep(10)
+
+        with patch.object(self.module, 'FETCH_TIMEOUT', 0.01):
+            with patch.object(self.module, 'download_avatar', side_effect=slow):
+                self.assertIsNone(await self.cache.get('aiocqhttp', '12345'))
+        with patch.object(self.module, 'download_avatar', new=AsyncMock(
+                return_value=b'x' * (self.module.MAX_BYTES + 1))):
+            self.assertIsNone(await self.cache.get('aiocqhttp', '54321'))
+        with patch.object(self.module, 'download_avatar',
+                          new=AsyncMock(return_value=b'avatar')):
+            for index in range(30):
+                await self.cache.get('aiocqhttp', str(100000 + index))
+        self.assertLessEqual(len(self.cache._cache), 16)
+
+    async def test_close_and_independent_caller_cancellation(self):
+        started, finish = asyncio.Event(), asyncio.Event()
+
+        async def download(user_id):
+            started.set()
+            await finish.wait()
+            return b'avatar'
+
+        with patch.object(self.module, 'download_avatar', side_effect=download):
+            first = asyncio.create_task(self.cache.get('aiocqhttp', '12345'))
+            await started.wait()
+            second = asyncio.create_task(self.cache.get('aiocqhttp', '12345'))
+            await asyncio.sleep(0)
+            first.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await first
+            finish.set()
+            self.assertEqual(await second, b'avatar')
+            started.clear()
+            finish.clear()
+            third = asyncio.create_task(self.cache.get('aiocqhttp', '54321'))
+            await started.wait()
+            await self.cache.close()
+            with self.assertRaises(asyncio.CancelledError):
+                await third
+            self.assertFalse(self.cache._pending)
+            self.assertFalse(self.cache._cache)
+
+    async def test_downloader_stream_limit_and_no_redirect(self):
+        class Body:
+            async def iter_chunked(self, size):
+                yield b'a' * (self_limit + 1)
+
+        self_limit = self.module.MAX_BYTES
+
+        class Response:
+            status = 200
+            content_length = None
+            content = Body()
+
+            def raise_for_status(self):
+                pass
+
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, *args):
+                pass
+
+        class Session(Response):
+            def get(self, url, **kwargs):
+                self_url.append((url, kwargs))
+                return Response()
+
+        self_url = []
+        with patch.object(self.module.aiohttp, 'ClientSession',
+                          return_value=Session()):
+            self.assertIsNone(await self.module.download_avatar('12345'))
+        self.assertEqual(self_url[0][0],
+                         'https://q1.qlogo.cn/g?b=qq&nk=12345&s=100')
+        self.assertFalse(self_url[0][1]['allow_redirects'])

tests/test_astrbot_package.py

@@ -7,7 +7,7 @@ import sys
 import tempfile
 from types import ModuleType, SimpleNamespace
 import unittest
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 
 ROOT = Path(__file__).resolve().parents[1]
@@ -115,7 +115,7 @@ class AstrBotPackageTests(unittest.IsolatedAsyncioTestCase):
                         return 'aiocqhttp'
 
                     def get_sender_id(self):
-                        return 'player'
+                        return '123456'
 
                     def get_sender_name(self):
                         return '离线测试'
@@ -139,14 +139,37 @@ class AstrBotPackageTests(unittest.IsolatedAsyncioTestCase):
                 again = await call('开通农场')
                 self.assertIsInstance(again, str)
                 self.assertIn('开通农场', await call('我的种子', 'bot-b'))
-                result = await call('我的农场 高清')
                 from io import BytesIO
                 from PIL import Image
+                avatar_buffer = BytesIO()
+                Image.new('RGB', (64, 64), '#ee2222').save(
+                    avatar_buffer, format='PNG',
+                )
+                avatar_bytes = avatar_buffer.getvalue()
+                with patch.object(plugin._avatars, 'get', new=AsyncMock(
+                        return_value=avatar_bytes)) as avatar_get:
+                    with patch.object(plugin.runtime, 'execute',
+                                      wraps=plugin.runtime.execute) as execute:
+                        result = await call('我的农场 高清')
+                        self.assertEqual(execute.call_args.kwargs['avatar'],
+                                         avatar_bytes)
+                        await call('农场详述')
+                        await call('地块编号')
+                        await call('我的种子 文字')
+                        self.assertIsNone(execute.call_args.kwargs['avatar'])
+                    self.assertEqual(avatar_get.await_count, 3)
+                    avatar_get.assert_awaited_with('aiocqhttp', '123456')
                 with Image.open(BytesIO(result[0])) as image:
                     self.assertEqual(image.width, 1536)
+                    red, green, blue = image.getpixel((70, 70))[:3]
+                    self.assertGreater(red, 180)
+                    self.assertLess(green, 90)
+                    self.assertLess(blue, 90)
                 self.assertIn('仅我的农场', await call('我的种子 高清'))
                 self.assertIn('暂无种子', await call('我的种子 文字'))
-                self.assertTrue((self.root / 'plugin-data/farm.sqlite3').is_file())
+                self.assertTrue(
+                    (self.root / 'plugin-data/farm.sqlite3').is_file(),
+                )
                 await plugin.terminate()
                 self.assertFalse(plugin.runtime.started)
             finally: