Skip to content

Commit f544e0f

Browse files
Fix lock property and add test coverage for RunUsage + RequestUsage
- Fix _lock property to check for '_RunUsage__lock' instead of '__lock' (hasattr with '__lock' string doesn't get name-mangled, so it always returned False) - Add test_run_usage_with_request_usage() to cover: * RunUsage.__add__(RequestUsage) - line 232->236 branch * RunUsage.incr(RequestUsage) - line 212->216 branch * Multiple accesses to _lock property - line 198->200 branch - Achieves 100% test coverage for usage.py
1 parent 6c4abb6 commit f544e0f

File tree

2 files changed

+29
-1
lines changed

2 files changed

+29
-1
lines changed

pydantic_ai_slim/pydantic_ai/usage.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,7 @@ class RunUsage(UsageBase):
195195
@property
196196
def _lock(self) -> threading.Lock:
197197
"""Get the lock, creating it if it doesn't exist."""
198-
if not hasattr(self, '__lock'):
198+
if not hasattr(self, '_RunUsage__lock'):
199199
self.__lock = threading.Lock()
200200
return self.__lock
201201

tests/test_usage_limits.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,34 @@ def test_add_usages():
248248
assert RunUsage() + RunUsage() == RunUsage()
249249

250250

251+
def test_run_usage_with_request_usage():
252+
"""Test RunUsage operations with RequestUsage to ensure coverage of RequestUsage branches."""
253+
run_usage = RunUsage(requests=1, input_tokens=10, output_tokens=20, tool_calls=1)
254+
request_usage = RequestUsage(input_tokens=5, output_tokens=10)
255+
256+
# Test __add__ with RequestUsage (covers line 232->236 branch)
257+
result = run_usage + request_usage
258+
assert result.requests == 2 # 1 + 1 (RequestUsage.requests property returns 1)
259+
assert result.input_tokens == 15
260+
assert result.output_tokens == 30
261+
assert result.tool_calls == 1 # RequestUsage doesn't have tool_calls
262+
263+
# Test incr with RequestUsage (covers line 212->216 branch)
264+
run_usage2 = RunUsage(requests=2, input_tokens=20, output_tokens=30, tool_calls=2)
265+
run_usage2.incr(request_usage)
266+
assert run_usage2.requests == 3 # 2 + 1
267+
assert run_usage2.input_tokens == 25 # 20 + 5
268+
assert run_usage2.output_tokens == 40 # 30 + 10
269+
assert run_usage2.tool_calls == 2 # Unchanged
270+
271+
# Test accessing _lock multiple times on same instance (covers line 198->200 branch)
272+
# First access creates the lock
273+
lock1 = run_usage._lock
274+
# Second access returns the existing lock (covers the hasattr return path)
275+
lock2 = run_usage._lock
276+
assert lock1 is lock2 # Same lock instance
277+
278+
251279
async def test_tool_call_limit() -> None:
252280
test_agent = Agent(TestModel())
253281

0 commit comments

Comments
 (0)