Skip to content

Commit

Permalink
Fix issue in async generator wrapper
Browse files Browse the repository at this point in the history
  • Loading branch information
TimPansino committed Jul 26, 2023
1 parent 5693dd2 commit c857358
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 2 deletions.
4 changes: 2 additions & 2 deletions newrelic/common/async_wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ async def wrapper(*args, **kwargs):
with trace:
while True:
try:
g.asend(value).send(None)
yielded = await g.asend(value)
except StopAsyncIteration as e:
# The underlying async generator has finished, return propagates a new StopAsyncIteration
return
Expand All @@ -107,7 +107,7 @@ async def wrapper(*args, **kwargs):
# An exception was thrown with .athrow(), propagate to the original async generator.
# Return value logic must be identical to .asend()
try:
g.athrow(type(e), e).send(None)
value = yield await g.athrow(type(e), e)
except StopAsyncIteration as e:
# The underlying async generator has finished, return propagates a new StopAsyncIteration
return
Expand Down
31 changes: 31 additions & 0 deletions tests/agent_features/_test_async_generator_trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,37 @@ async def _test():
assert full_metrics[key].total_exclusive_call_time < 0.2


@validate_transaction_metrics(
"test_asend_receives_a_value",
background_task=True,
scoped_metrics=[("Function/agen", 1)],
rollup_metrics=[("Function/agen", 1)],
)
def test_asend_receives_a_value(event_loop):
_received = []
@function_trace(name="agen")
async def agen():
value = yield
_received.append(value)
yield value

@background_task(name="test_asend_receives_a_value")
async def _test():
gen = agen()

# kickstart the coroutine
await anext(gen)

assert await gen.asend("foobar") == "foobar"
assert _received and _received[0] == "foobar"

# finish consumption of the coroutine if necessary
async for _ in gen:
pass

event_loop.run_until_complete(_test())


@validate_transaction_metrics(
"test_athrow_yields_a_value",
background_task=True,
Expand Down

0 comments on commit c857358

Please sign in to comment.