Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
1.9 kB · 64 lines
Python
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465from __future__ import annotations
import asyncioimport importlibimport unittestfrom unittest.mock import AsyncMock, patch
rlm_module = importlib.import_module("rlm")
class RlmContextUsageTest(unittest.TestCase): def test_get_context_usage_success(self) -> None: host_request = AsyncMock( return_value={ "usage": { "tokens": 15000, "context_window": 128000, "percent": 11.71875, } } )
with patch.object(rlm_module, "host_request", host_request): usage = asyncio.run(rlm_module.rlm.get_context_usage())
self.assertIsNotNone(usage) self.assertEqual(usage.tokens, 15000) self.assertEqual(usage.context_window, 128000) self.assertAlmostEqual(usage.percent, 11.71875)
def test_get_context_usage_null_tokens_after_compaction(self) -> None: host_request = AsyncMock( return_value={ "usage": { "tokens": None, "context_window": 128000, "percent": None, } } )
with patch.object(rlm_module, "host_request", host_request): usage = asyncio.run(rlm_module.get_context_usage())
self.assertIsNotNone(usage) self.assertIsNone(usage.tokens) self.assertEqual(usage.context_window, 128000) self.assertIsNone(usage.percent)
def test_get_context_usage_none_when_no_usage(self) -> None: host_request = AsyncMock( return_value={ "usage": None, } )
with patch.object(rlm_module, "host_request", host_request): usage = asyncio.run(rlm_module.rlm.context_usage())
self.assertIsNone(usage)
if __name__ == "__main__": unittest.main()