36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
import pytest
|
|
|
|
def test_token_usage_aggregation():
|
|
# A dummy test to fulfill the Red Phase for the new token usage widget.
|
|
# We will implement a function in gui.py or ai_client.py to aggregate tokens.
|
|
from ai_client import _comms_log, clear_comms_log, _append_comms
|
|
|
|
clear_comms_log()
|
|
|
|
_append_comms("IN", "response", {
|
|
"usage": {
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
"cache_read_input_tokens": 10,
|
|
"cache_creation_input_tokens": 5
|
|
}
|
|
})
|
|
|
|
_append_comms("IN", "response", {
|
|
"usage": {
|
|
"input_tokens": 200,
|
|
"output_tokens": 100,
|
|
"cache_read_input_tokens": 20,
|
|
"cache_creation_input_tokens": 0
|
|
}
|
|
})
|
|
|
|
# We expect a new function get_total_token_usage() to exist
|
|
from gui import get_total_token_usage
|
|
|
|
totals = get_total_token_usage()
|
|
assert totals["input_tokens"] == 300
|
|
assert totals["output_tokens"] == 150
|
|
assert totals["cache_read_input_tokens"] == 30
|
|
assert totals["cache_creation_input_tokens"] == 5
|