add mock LM for tests, add test for semaphore

Author: alexzhang13

rlm/clients/mock_lm.py

@@ -0,0 +1,68 @@
+"""
+Mock LM client for testing without a real model.
+
+Use this to test LMHandler, REPL integrations, or any code that needs a BaseLM
+without calling an external API.
+"""
+
+from typing import Any, Callable
+
+from rlm.clients.base_lm import BaseLM
+from rlm.core.types import ModelUsageSummary, UsageSummary
+
+
+class MockLM(BaseLM):
+    """
+    In-memory mock LM that implements BaseLM for testing.
+
+    - completion() / acompletion(): return a simple echo-style response by default,
+      or use the optional response list/callable if provided.
+    - get_usage_summary() / get_last_usage(): return fixed usage for tests.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "mock-model",
+        responses: list[str] | None = None,
+        response_fn: Callable[[str | dict[str, Any]], str] | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(model_name=model_name, **kwargs)
+        self._responses = list(responses) if responses is not None else None
+        self._response_fn = response_fn
+        self._call_count = 0
+        self._default_usage = ModelUsageSummary(
+            total_calls=0, total_input_tokens=0, total_output_tokens=0
+        )
+
+    def completion(self, prompt: str | dict[str, Any]) -> str:
+        self._call_count += 1
+        if self._responses is not None:
+            if not self._responses:
+                raise IndexError("MockLM: no more responses in list")
+            return self._responses.pop(0)
+        if self._response_fn is not None:
+            return self._response_fn(prompt)
+        prompt_str = prompt if isinstance(prompt, str) else str(prompt)[:80]
+        return f"Mock response to: {prompt_str}"
+
+    async def acompletion(self, prompt: str | dict[str, Any]) -> str:
+        return self.completion(prompt)
+
+    def get_usage_summary(self) -> UsageSummary:
+        return UsageSummary(
+            model_usage_summaries={
+                self.model_name: ModelUsageSummary(
+                    total_calls=self._call_count,
+                    total_input_tokens=self._call_count * 10,
+                    total_output_tokens=self._call_count * 10,
+                )
+            }
+        )
+
+    def get_last_usage(self) -> ModelUsageSummary:
+        return ModelUsageSummary(
+            total_calls=1,
+            total_input_tokens=10,
+            total_output_tokens=10,
+        )

tests/test_lm_handler.py

@@ -0,0 +1,47 @@
+"""Tests for LMHandler using MockLM (no real LM required)."""
+
+import pytest
+
+from rlm.clients.mock_lm import MockLM
+from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched
+from rlm.core.lm_handler import LMHandler
+
+
+def test_lm_handler_single_request():
+    """Single prompt request returns success and echo-style content."""
+    mock = MockLM(responses=["hello back"])
+    with LMHandler(client=mock) as handler:
+        request = LMRequest(prompt="hello")
+        response = send_lm_request(handler.address, request)
+    assert response.success
+    assert response.chat_completion is not None
+    assert response.chat_completion.response == "hello back"
+
+
+def test_lm_handler_batched_request():
+    """Batched prompts return one response per prompt in order."""
+    responses = [f"r{i}" for i in range(5)]
+    mock = MockLM(responses=responses)
+    with LMHandler(client=mock, batch_max_concurrent=3) as handler:
+        prompts = [f"prompt-{i}" for i in range(5)]
+        result = send_lm_request_batched(handler.address, prompts)
+    assert len(result) == 5
+    for i, resp in enumerate(result):
+        assert resp.success, resp.error
+        assert resp.chat_completion is not None
+        assert resp.chat_completion.response == f"r{i}"
+
+
+def test_lm_handler_batched_many_prompts_semaphore_cap():
+    """Many prompts complete successfully with semaphore limiting concurrency."""
+    # 50 prompts, max 4 concurrent: should still all complete
+    count = 50
+    responses = [f"resp-{i}" for i in range(count)]
+    mock = MockLM(responses=responses)
+    with LMHandler(client=mock, batch_max_concurrent=4) as handler:
+        prompts = [f"p-{i}" for i in range(count)]
+        result = send_lm_request_batched(handler.address, prompts)
+    assert len(result) == count
+    for i, resp in enumerate(result):
+        assert resp.success, (i, resp.error)
+        assert resp.chat_completion.response == f"resp-{i}"