mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-07-17 19:27:11 +08:00
ad64e17a34
Two independent robustness fixes: 1. ChatChannel.cancel_session / cancel_all_session raised KeyError when a session existed in self.sessions but no future had been dispatched yet. self.sessions[sid] is created in produce(), but self.futures[sid] is only created later in consume() on first dispatch. Cancelling in that window (e.g. user sends a message then immediately cancels) crashed the cancel path. Use self.futures.get(sid, []) so an absent entry is a no-op. 2. compress_imgfile decremented JPEG quality by 5 with no lower bound. For an image that cannot be compressed below max_size, quality went 0, negative, ... — the loop never terminated and passed invalid quality values to PIL. Add a min_quality floor (10) and return the best effort once reached. Adds tests/test_robustness_fixes.py covering both paths (5 tests).
122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
# encoding:utf-8
|
|
"""
|
|
Unit tests for robustness fixes:
|
|
1. ChatChannel.cancel_session / cancel_all_session must not raise KeyError
|
|
when a session has been produced but no task has been dispatched yet
|
|
(so self.futures[session_id] does not exist).
|
|
2. common.utils.compress_imgfile must terminate (no infinite loop / invalid
|
|
PIL quality) when an image cannot be compressed below max_size.
|
|
"""
|
|
import io
|
|
import os
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from unittest.mock import MagicMock
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
|
|
|
|
# =============================================================================
|
|
# 1. cancel_session / cancel_all_session KeyError regression
|
|
# =============================================================================
|
|
|
|
class TestCancelSessionMissingFutures(unittest.TestCase):
|
|
"""A session may exist in self.sessions before any future is recorded."""
|
|
|
|
def _make_channel(self):
|
|
# Import lazily and build a bare object without running __init__,
|
|
# to avoid pulling the full channel setup / config.
|
|
from channel.chat_channel import ChatChannel
|
|
|
|
ch = ChatChannel.__new__(ChatChannel)
|
|
import threading
|
|
|
|
ch.lock = threading.RLock()
|
|
# A produced session whose future has NOT been dispatched yet.
|
|
queue = MagicMock()
|
|
queue.qsize.return_value = 0
|
|
semaphore = MagicMock()
|
|
ch.sessions = {"sid": [queue, semaphore]}
|
|
ch.futures = {} # intentionally empty: consume() never ran
|
|
return ch
|
|
|
|
def test_cancel_session_no_futures_entry(self):
|
|
ch = self._make_channel()
|
|
# Should not raise KeyError.
|
|
try:
|
|
ch.cancel_session("sid")
|
|
except KeyError:
|
|
self.fail("cancel_session raised KeyError when futures entry missing")
|
|
|
|
def test_cancel_all_session_no_futures_entry(self):
|
|
ch = self._make_channel()
|
|
try:
|
|
ch.cancel_all_session()
|
|
except KeyError:
|
|
self.fail("cancel_all_session raised KeyError when futures entry missing")
|
|
|
|
def test_cancel_session_cancels_existing_futures(self):
|
|
ch = self._make_channel()
|
|
fut = MagicMock()
|
|
ch.futures["sid"] = [fut]
|
|
ch.cancel_session("sid")
|
|
fut.cancel.assert_called_once()
|
|
|
|
|
|
# =============================================================================
|
|
# 2. compress_imgfile termination
|
|
# =============================================================================
|
|
|
|
class TestCompressImgfileTermination(unittest.TestCase):
|
|
"""compress_imgfile must always return, even for incompressible input."""
|
|
|
|
def setUp(self):
|
|
# Skip if Pillow is not available in the test environment.
|
|
try:
|
|
import PIL # noqa: F401
|
|
except ImportError:
|
|
self.skipTest("Pillow not installed")
|
|
|
|
def _make_image_buf(self, size=(64, 64)):
|
|
from PIL import Image
|
|
import random
|
|
|
|
img = Image.new("RGB", size)
|
|
# Fill with random noise so JPEG cannot compress it well.
|
|
pixels = [
|
|
(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
|
|
for _ in range(size[0] * size[1])
|
|
]
|
|
img.putdata(pixels)
|
|
buf = io.BytesIO()
|
|
img.save(buf, "JPEG", quality=95)
|
|
buf.seek(0)
|
|
return buf
|
|
|
|
def test_returns_when_target_unreachable(self):
|
|
from common.utils import compress_imgfile
|
|
|
|
buf = self._make_image_buf()
|
|
# An impossibly small target that even quality=10 won't reach.
|
|
out = compress_imgfile(buf, max_size=10)
|
|
self.assertIsInstance(out, io.BytesIO)
|
|
# Verify the result is still a valid JPEG (PIL never got invalid quality).
|
|
from PIL import Image
|
|
|
|
out.seek(0)
|
|
img = Image.open(out)
|
|
img.verify()
|
|
|
|
def test_no_compression_needed_returns_same_object(self):
|
|
from common.utils import compress_imgfile
|
|
|
|
buf = self._make_image_buf()
|
|
size = buf.getbuffer().nbytes
|
|
out = compress_imgfile(buf, max_size=size + 1)
|
|
self.assertIs(out, buf)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|