From fedc9b6bed2b0302bb1791b682aa570982fe91da Mon Sep 17 00:00:00 2001 From: taoyifan Date: Wed, 31 Dec 2025 16:42:45 +0800 Subject: [PATCH] Fix decorator exception bug. Change-Id: I9c389fe66d63f8e65233595f5c354d97957b03a1 --- cozeloop/decorator/decorator.py | 12 ++-- cozeloop/internal/version.py | 2 +- tests/decorator/test_decorator.py | 113 ++++++++++++++++++++++++++++++ 3 files changed, 121 insertions(+), 6 deletions(-) diff --git a/cozeloop/decorator/decorator.py b/cozeloop/decorator/decorator.py index 9d5106d..de9772a 100644 --- a/cozeloop/decorator/decorator.py +++ b/cozeloop/decorator/decorator.py @@ -101,7 +101,8 @@ def sync_wrapper(*args: Any, **kwargs: Any): span.set_tags(tags) span.finish() - return res + if res is not None: + return res @wraps(func) async def async_wrapper(*args: Any, **kwargs: Any): @@ -139,7 +140,8 @@ async def async_wrapper(*args: Any, **kwargs: Any): span.set_tags(tags) span.finish() - return res + if res is not None: + return res @wraps(func) def gen_wrapper(*args: Any, **kwargs: Any): @@ -247,7 +249,7 @@ def sync_stream_wrapper(*args: Any, **kwargs: Any): span.set_input(input) span.set_tags(tags) - if not hasattr(res, "__iter__"): + if not hasattr(res, "__iter__") and res is not None: return res @wraps(func) @@ -286,7 +288,7 @@ async def async_stream_wrapper(*args: Any, **kwargs: Any): span.set_input(input) span.set_tags(tags) - if not hasattr(res, "__aiter__"): + if not hasattr(res, "__aiter__") and res is not None: return res if is_async_gen_func(func): @@ -421,7 +423,7 @@ async def __aiter__(self) -> AsyncIterator[S]: await self._aend() raise except Exception as e: - await self._aend() + await self._aend(e) raise e else: await self._aend() diff --git a/cozeloop/internal/version.py b/cozeloop/internal/version.py index 6d61e34..3fa0c40 100644 --- a/cozeloop/internal/version.py +++ b/cozeloop/internal/version.py @@ -1,4 +1,4 @@ # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT -VERSION = 'v0.1.22' +VERSION = 'v0.1.23' diff --git a/tests/decorator/test_decorator.py b/tests/decorator/test_decorator.py index 65822b6..250e642 100644 --- a/tests/decorator/test_decorator.py +++ b/tests/decorator/test_decorator.py @@ -368,3 +368,116 @@ async def collect(): assert result == [1, 2, 3] assert span.output == '[2, 4, 6]' assert span.finished is True + + +class TestCozeLoopDecoratorExceptionHandling: + """测试 CozeLoopDecorator 是否正确处理并重新抛出异常，而不是将其掩盖""" + + def test_sync_func_exception(self, monkeypatch): + span = SpanMock() + monkeypatch.setattr(decorator_module, "start_span", lambda *args, **kwargs: span) + + @decorator_module.CozeLoopDecorator().observe() + def risky_func(): + raise ValueError("Sync error") + + with pytest.raises(ValueError, match="Sync error"): + risky_func() + + assert span.error == "Sync error" + assert span.finished is True + + def test_async_func_exception(self, monkeypatch): + span = SpanMock() + monkeypatch.setattr(decorator_module, "start_span", lambda *args, **kwargs: span) + + @decorator_module.CozeLoopDecorator().observe() + async def risky_async_func(): + raise ValueError("Async error") + + with pytest.raises(ValueError, match="Async error"): + asyncio.run(risky_async_func()) + + assert span.error == "Async error" + assert span.finished is True + + def test_generator_func_exception(self, monkeypatch): + span = SpanMock() + monkeypatch.setattr(decorator_module, "start_span", lambda *args, **kwargs: span) + + @decorator_module.CozeLoopDecorator().observe() + def risky_gen(): + yield 1 + raise ValueError("Gen error") + + gen = risky_gen() + assert next(gen) == 1 + with pytest.raises(ValueError, match="Gen error"): + next(gen) + + assert span.error == "Gen error" + assert span.finished is True + + def test_async_generator_func_exception(self, monkeypatch): + span = SpanMock() + monkeypatch.setattr(decorator_module, "start_span", lambda *args, **kwargs: span) + + @decorator_module.CozeLoopDecorator().observe() + async def risky_agen(): + yield 1 + raise ValueError("Async gen error") + + async def collect(): + agen = risky_agen() + assert (await agen.__anext__()) == 1 + with pytest.raises(ValueError, match="Async gen error"): + await agen.__anext__() + + asyncio.run(collect()) + + assert span.error == "Async gen error" + assert span.finished is True + + def test_sync_stream_wrapper_exception(self, monkeypatch): + span = SpanMock() + monkeypatch.setattr(decorator_module, "start_span", lambda *args, **kwargs: span) + + def risky_iter_func(): + class RiskyIter: + def __iter__(self): + return self + def __next__(self): + raise ValueError("Stream error") + return RiskyIter() + + decorated = decorator_module.CozeLoopDecorator().observe(process_iterator_outputs=lambda x: x)(risky_iter_func) + stream = decorated() + with pytest.raises(ValueError, match="Stream error"): + list(stream) + + assert span.error == "Stream error" + assert span.finished is True + + def test_async_stream_wrapper_exception(self, monkeypatch): + span = SpanMock() + monkeypatch.setattr(decorator_module, "start_span", lambda *args, **kwargs: span) + + async def risky_aiter_func(): + class RiskyAIter: + def __aiter__(self): + return self + async def __anext__(self): + raise ValueError("Async stream error") + return RiskyAIter() + + async def collect(): + decorated = decorator_module.CozeLoopDecorator().observe(process_iterator_outputs=lambda x: x)(risky_aiter_func) + stream = await decorated() + with pytest.raises(ValueError, match="Async stream error"): + async for _ in stream: + pass + + asyncio.run(collect()) + + assert span.error == "Async stream error" + assert span.finished is True