Repository: oumi-ai/oumi
File: tests/unit/inference/test_remote_inference_engine.py (Python)

Code before:
        m.post(
            _TARGET_SERVER,
            status=200,
            body="Invalid JSON content",
            content_type="text/plain",
        )

        engine = RemoteInferenceEngine(
            model_params=_get_default_model_params(),
            remote_params=RemoteParams(
                api_url=_TARGET_SERVER,
                max_retries=2,
                retry_backoff_base=0.1,  # Small value for testing
                retry_backoff_max=0.3,
            ),
        )
        conversation = Conversation(
            messages=[
                Message(
                    role=Role.USER,
                    content="Hello world!",
                ),
            ],
        )
        
        with pytest.raises(RuntimeError) as exc_info:
            engine.infer_online(
                [conversation],
                _get_default_inference_config(),
            )
        
        assert "Failed to parse response as JSON" in str(exc_info.value)
        assert "Response content type: text/plain" in str(exc_info.value)


def test_infer_online_exponential_backoff():
    """Test that the engine implements exponential backoff correctly."""
    request_times = []
    
    def callback(url, **kwargs):
        request_times.append(time.time())
        # Fail until the last attempt
        if len(request_times) < 3:
            return CallbackResult(status=500, body=json.dumps({"error": {"message": "Server Error"}}))
        return CallbackResult(
            status=200,
            payload={
                "choices": [
                    {
                        "message": {
                            "role": "assistant",

Review comment (suggestion):
You can use `pytest.raises(RuntimeError, match="...")` here instead of the individual asserts below.

Code after:
        inference_config = _get_default_inference_config()
        inference_config.remote_params = engine._remote_params
        result = engine.infer_online(
            [conversation],
            inference_config,
        )
        assert len(result) == 1
        assert result[0].messages[-1].content == "The first time I saw"
        assert result[0].messages[-1].role == Role.ASSISTANT


def test_infer_online_handles_invalid_content():
    """Test that the engine properly handles invalid content responses."""
    with aioresponses() as m:
        m.post(
            _TARGET_SERVER,
            status=200,
            body=json.dumps({"error": {"message": "Invalid JSON content"}}),
            content_type="application/json",
        )

        engine = RemoteInferenceEngine(
            model_params=_get_default_model_params(),
            remote_params=RemoteParams(
                api_url=_TARGET_SERVER,
                max_retries=2,
                retry_backoff_base=0.1,  # Small value for testing
                retry_backoff_max=0.3,
            ),
        )
        conversation = Conversation(
            messages=[
                Message(
                    role=Role.USER,
                    content="Hello world!",
                ),
            ],
        )

        inference_config = _get_default_inference_config()
        inference_config.remote_params = engine._remote_params

        async def mock_sleep(delay):
            pass

        with patch("asyncio.sleep", side_effect=mock_sleep):
            with pytest.raises(RuntimeError, match="API error: Invalid JSON content"):
                engine.infer_online(
                    [conversation],
                    inference_config,
                )