Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions cozeloop/decorator/decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion cozeloop/internal/version.py
Original file line number Diff line number Diff line change
@@ -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'
113 changes: 113 additions & 0 deletions tests/decorator/test_decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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