from __future__ import annotations

import asyncio
import importlib
import unittest
from 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()