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

Code before:
        m.post(_TARGET_SERVER, callback=callback)

        engine = RemoteInferenceEngine(
            model_params=_get_default_model_params(),
            remote_params=RemoteParams(
                api_url=_TARGET_SERVER,
                max_retries=2,
                retry_backoff_base=0.2,  # Small values for testing
                retry_backoff_max=1.0,
            ),
        )
        conversation = Conversation(
            messages=[Message(role=Role.USER, content="Hello")],
        )
        
        result = engine.infer_online([conversation], _get_default_inference_config())
        
        # Verify the result
        assert len(result) == 1
        assert result[0].messages[-1].content == "Success after retries"
        
        # Verify timing between requests
        delays = [request_times[i] - request_times[i-1] for i in range(1, len(request_times))]
        # First retry should wait ~0.2s, second retry ~0.4s
        assert 0.15 <= delays[0] <= 0.25  # Allow some timing variance
        assert 0.35 <= delays[1] <= 0.45  # Allow some timing variance


Review comment (bug):
A safer way to test this would be to mock out `asyncio.sleep()`. You can then just verify if `asyncio.sleep` was called a certain number of times with specific arguments. This also helps our unit tests run faster without having to actually sleep :)

See https://stackoverflow.com/questions/57747998/how-to-make-unit-test-with-asyncio-sleep-contained-code for some examples of mocking this method.

Code after:

    def callback(url, **kwargs):
        # Fail until the last attempt
        if len(sleep_calls) < 3:
            return CallbackResult(
                status=500,
                body=json.dumps({"error": {"message": "Server Error"}}),
                content_type="application/json",
            )
        return CallbackResult(
            status=200,
            body=json.dumps(
                {
                    "choices": [
                        {
                            "message": {
                                "role": "assistant",
                                "content": "Success after retries",
                            }
                        }
                    ]
                }
            ),
            content_type="application/json",
        )

    with aioresponses() as m:
        m.post(_TARGET_SERVER, callback=callback, repeat=True)

        with patch("asyncio.sleep", side_effect=mock_sleep):
            engine = RemoteInferenceEngine(
                model_params=_get_default_model_params(),
                remote_params=RemoteParams(
                    api_url=_TARGET_SERVER,
                    max_retries=2,
                    retry_backoff_base=0.2,  # Small values for testing
                    retry_backoff_max=1.0,
                ),
            )
            conversation = Conversation(
                messages=[Message(role=Role.USER, content="Hello")],
            )

            remote_params = RemoteParams(
                api_url=_TARGET_SERVER,
                max_retries=2,
                retry_backoff_base=0.2,
                retry_backoff_max=1.0,
            )
            inference_config = _get_default_inference_config()
            inference_config.remote_params = remote_params