mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-07-17 03:03:19 +08:00
fix: avoid KeyError on cancel and infinite loop in image compression
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).
This commit is contained in:
@@ -519,7 +519,10 @@ class ChatChannel(Channel):
|
||||
def cancel_session(self, session_id):
|
||||
with self.lock:
|
||||
if session_id in self.sessions:
|
||||
for future in self.futures[session_id]:
|
||||
# futures[session_id] is only created in consume() when a task is
|
||||
# dispatched, so it may be absent if cancel happens right after
|
||||
# produce() but before the first dispatch. Default to [].
|
||||
for future in self.futures.get(session_id, []):
|
||||
future.cancel()
|
||||
cnt = self.sessions[session_id][0].qsize()
|
||||
if cnt > 0:
|
||||
@@ -529,7 +532,7 @@ class ChatChannel(Channel):
|
||||
def cancel_all_session(self):
|
||||
with self.lock:
|
||||
for session_id in self.sessions:
|
||||
for future in self.futures[session_id]:
|
||||
for future in self.futures.get(session_id, []):
|
||||
future.cancel()
|
||||
cnt = self.sessions[session_id][0].qsize()
|
||||
if cnt > 0:
|
||||
|
||||
@@ -27,10 +27,14 @@ def compress_imgfile(file, max_size):
|
||||
img = Image.open(file)
|
||||
rgb_image = img.convert("RGB")
|
||||
quality = 95
|
||||
min_quality = 10
|
||||
while True:
|
||||
out_buf = io.BytesIO()
|
||||
rgb_image.save(out_buf, "JPEG", quality=quality)
|
||||
if fsize(out_buf) <= max_size:
|
||||
if fsize(out_buf) <= max_size or quality <= min_quality:
|
||||
# Stop at min_quality: further decrements would pass an invalid
|
||||
# quality (<1) to PIL and the loop would otherwise never terminate
|
||||
# for images that cannot be compressed below max_size.
|
||||
return out_buf
|
||||
quality -= 5
|
||||
|
||||
|
||||
121
tests/test_robustness_fixes.py
Normal file
121
tests/test_robustness_fixes.py
Normal file
@@ -0,0 +1,121 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user