Skip to content

Commit b5f510e

Browse files
sjmonsonsjmonson
committed
Add tests for openai completion code
1 parent 1b2c185 commit b5f510e

File tree

1 file changed

+305
-1
lines changed

1 file changed

+305
-1
lines changed

‎tests/unit/backend/test_openai_backend.py

Lines changed: 305 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
import base64
88
from pathlib import Path
9-
from unittest.mock import Mock, patch
9+
from unittest.mock import AsyncMock, Mock, patch
1010

1111
import httpx
1212
import pytest
@@ -845,3 +845,307 @@ async def test_chat_completions_content_formatting(self):
845845
assert body["messages"] == expected_messages
846846
finally:
847847
await backend.process_shutdown()
848+
849+
@pytest.mark.regression
850+
@pytest.mark.asyncio
851+
async def test_openai_backend_validate_no_models_available(self):
852+
"""Test validate method when no models are available.
853+
854+
### WRITTEN BY AI ###
855+
"""
856+
backend = OpenAIHTTPBackend(target="http://test")
857+
await backend.process_startup()
858+
859+
try:
860+
# Mock endpoints to fail, then available_models to return empty list
861+
def mock_get_fail(*args, **kwargs):
862+
raise httpx.HTTPStatusError("Error", request=Mock(), response=Mock())
863+
864+
with (
865+
patch.object(backend._async_client, "get", side_effect=mock_get_fail),
866+
patch.object(backend, "available_models", return_value=[]),
867+
patch.object(backend, "text_completions", side_effect=mock_get_fail),
868+
pytest.raises(
869+
RuntimeError,
870+
match="No model available and could not set a default model",
871+
),
872+
):
873+
await backend.validate()
874+
finally:
875+
await backend.process_shutdown()
876+
877+
@pytest.mark.sanity
878+
@pytest.mark.asyncio
879+
async def test_text_completions_streaming(self):
880+
"""Test text_completions with streaming enabled."""
881+
backend = OpenAIHTTPBackend(target="http://test", model="gpt-4")
882+
await backend.process_startup()
883+
884+
try:
885+
# Mock streaming response
886+
mock_stream = Mock()
887+
mock_stream.raise_for_status = Mock()
888+
889+
async def mock_aiter_lines():
890+
lines = [
891+
'data: {"choices":[{"text":"Hello"}], "usage":{"prompt_tokens":5,"completion_tokens":1}}', # noqa: E501
892+
'data: {"choices":[{"text":" world"}], "usage":{"prompt_tokens":5,"completion_tokens":2}}', # noqa: E501
893+
'data: {"choices":[{"text":"!"}], "usage":{"prompt_tokens":5,"completion_tokens":3}}', # noqa: E501
894+
"data: [DONE]",
895+
]
896+
for line in lines:
897+
yield line
898+
899+
mock_stream.aiter_lines = mock_aiter_lines
900+
901+
mock_client_stream = AsyncMock()
902+
mock_client_stream.__aenter__ = AsyncMock(return_value=mock_stream)
903+
mock_client_stream.__aexit__ = AsyncMock(return_value=None)
904+
905+
with patch.object(
906+
backend._async_client, "stream", return_value=mock_client_stream
907+
):
908+
results = []
909+
async for result in backend.text_completions(
910+
prompt="test prompt", request_id="req-123", stream_response=True
911+
):
912+
results.append(result)
913+
914+
# Should get initial None, then tokens, then final with usage
915+
assert len(results) >= 3
916+
assert results[0] == (None, None) # Initial yield
917+
assert all(
918+
isinstance(result[0], str) for result in results[1:]
919+
) # Has text content
920+
assert all(
921+
isinstance(result[1], UsageStats) for result in results[1:]
922+
) # Has usage stats
923+
assert all(
924+
result[1].output_tokens == i for i, result in enumerate(results[1:], 1)
925+
)
926+
finally:
927+
await backend.process_shutdown()
928+
929+
@pytest.mark.sanity
930+
@pytest.mark.asyncio
931+
async def test_chat_completions_streaming(self):
932+
"""Test chat_completions with streaming enabled.
933+
934+
### WRITTEN BY AI ###
935+
"""
936+
backend = OpenAIHTTPBackend(target="http://test", model="gpt-4")
937+
await backend.process_startup()
938+
939+
try:
940+
# Mock streaming response
941+
mock_stream = Mock()
942+
mock_stream.raise_for_status = Mock()
943+
944+
async def mock_aiter_lines():
945+
lines = [
946+
'data: {"choices":[{"delta":{"content":"Hi"}}]}',
947+
'data: {"choices":[{"delta":{"content":" there"}}]}',
948+
'data: {"choices":[{"delta":{"content":"!"}}]}',
949+
'data: {"usage":{"prompt_tokens":3,"completion_tokens":3}}',
950+
"data: [DONE]",
951+
]
952+
for line in lines:
953+
yield line
954+
955+
mock_stream.aiter_lines = mock_aiter_lines
956+
957+
mock_client_stream = AsyncMock()
958+
mock_client_stream.__aenter__ = AsyncMock(return_value=mock_stream)
959+
mock_client_stream.__aexit__ = AsyncMock(return_value=None)
960+
961+
with patch.object(
962+
backend._async_client, "stream", return_value=mock_client_stream
963+
):
964+
results = []
965+
async for result in backend.chat_completions(
966+
content="Hello", request_id="req-456", stream_response=True
967+
):
968+
results.append(result)
969+
970+
# Should get initial None, then deltas, then final with usage
971+
assert len(results) >= 3
972+
assert results[0] == (None, None) # Initial yield
973+
assert any(result[0] for result in results if result[0]) # Has content
974+
assert any(result[1] for result in results if result[1]) # Has usage stats
975+
finally:
976+
await backend.process_shutdown()
977+
978+
@pytest.mark.regression
979+
@pytest.mark.asyncio
980+
async def test_streaming_response_edge_cases(self):
981+
"""Test streaming response edge cases for line processing.
982+
983+
### WRITTEN BY AI ###
984+
"""
985+
backend = OpenAIHTTPBackend(target="http://test", model="gpt-4")
986+
await backend.process_startup()
987+
988+
try:
989+
# Mock streaming response with edge cases
990+
mock_stream = Mock()
991+
mock_stream.raise_for_status = Mock()
992+
993+
async def mock_aiter_lines():
994+
lines = [
995+
"", # Empty line
996+
" ", # Whitespace only
997+
"not data line", # Line without data prefix
998+
'data: {"choices":[{"text":"Hello"}]}', # Valid data
999+
"data: [DONE]", # End marker
1000+
]
1001+
for line in lines:
1002+
yield line
1003+
1004+
mock_stream.aiter_lines = mock_aiter_lines
1005+
1006+
mock_client_stream = AsyncMock()
1007+
mock_client_stream.__aenter__ = AsyncMock(return_value=mock_stream)
1008+
mock_client_stream.__aexit__ = AsyncMock(return_value=None)
1009+
1010+
with patch.object(
1011+
backend._async_client, "stream", return_value=mock_client_stream
1012+
):
1013+
results = []
1014+
async for result in backend.text_completions(
1015+
prompt="test", request_id="req-123", stream_response=True
1016+
):
1017+
results.append(result)
1018+
1019+
# Should get initial None and the valid response
1020+
assert len(results) == 2
1021+
assert results[0] == (None, None)
1022+
assert results[1][0] == "Hello"
1023+
finally:
1024+
await backend.process_shutdown()
1025+
1026+
@pytest.mark.sanity
1027+
def test_openai_backend_get_chat_message_media_item_jpeg_file(self):
1028+
"""Test _get_chat_message_media_item with JPEG file path.
1029+
1030+
### WRITTEN BY AI ###
1031+
"""
1032+
backend = OpenAIHTTPBackend(target="http://test")
1033+
1034+
# Create a mock Path object for JPEG file
1035+
mock_jpeg_path = Mock(spec=Path)
1036+
mock_jpeg_path.suffix.lower.return_value = ".jpg"
1037+
1038+
# Mock Image.open to return a mock image
1039+
mock_image = Mock(spec=Image.Image)
1040+
mock_image.tobytes.return_value = b"fake_jpeg_data"
1041+
1042+
with patch("guidellm.backend.openai.Image.open", return_value=mock_image):
1043+
result = backend._get_chat_message_media_item(mock_jpeg_path)
1044+
1045+
expected_data = base64.b64encode(b"fake_jpeg_data").decode("utf-8")
1046+
expected = {
1047+
"type": "image",
1048+
"image": {"url": f"data:image/jpeg;base64,{expected_data}"},
1049+
}
1050+
assert result == expected
1051+
1052+
@pytest.mark.sanity
1053+
def test_openai_backend_get_chat_message_media_item_wav_file(self):
1054+
"""Test _get_chat_message_media_item with WAV file path.
1055+
1056+
### WRITTEN BY AI ###
1057+
"""
1058+
backend = OpenAIHTTPBackend(target="http://test")
1059+
1060+
# Create a mock Path object for WAV file
1061+
mock_wav_path = Mock(spec=Path)
1062+
mock_wav_path.suffix.lower.return_value = ".wav"
1063+
mock_wav_path.read_bytes.return_value = b"fake_wav_data"
1064+
1065+
result = backend._get_chat_message_media_item(mock_wav_path)
1066+
1067+
expected_data = base64.b64encode(b"fake_wav_data").decode("utf-8")
1068+
expected = {
1069+
"type": "input_audio",
1070+
"input_audio": {"data": expected_data, "format": "wav"},
1071+
}
1072+
assert result == expected
1073+
1074+
@pytest.mark.sanity
1075+
def test_openai_backend_get_chat_messages_with_pil_image(self):
1076+
"""Test _get_chat_messages with PIL Image in content list.
1077+
1078+
### WRITTEN BY AI ###
1079+
"""
1080+
backend = OpenAIHTTPBackend(target="http://test")
1081+
1082+
# Create a mock PIL Image
1083+
mock_image = Mock(spec=Image.Image)
1084+
mock_image.tobytes.return_value = b"fake_image_bytes"
1085+
1086+
content = ["Hello", mock_image, "world"]
1087+
1088+
result = backend._get_chat_messages(content)
1089+
1090+
# Should have one user message with mixed content
1091+
assert len(result) == 1
1092+
assert result[0]["role"] == "user"
1093+
assert len(result[0]["content"]) == 3
1094+
1095+
# Check text items
1096+
assert result[0]["content"][0] == {"type": "text", "text": "Hello"}
1097+
assert result[0]["content"][2] == {"type": "text", "text": "world"}
1098+
1099+
# Check image item
1100+
image_item = result[0]["content"][1]
1101+
assert image_item["type"] == "image"
1102+
assert "data:image/jpeg;base64," in image_item["image"]["url"]
1103+
1104+
@pytest.mark.regression
1105+
@pytest.mark.asyncio
1106+
async def test_resolve_timing_edge_cases(self):
1107+
"""Test resolve method timing edge cases.
1108+
1109+
### WRITTEN BY AI ###
1110+
"""
1111+
backend = OpenAIHTTPBackend(target="http://test")
1112+
await backend.process_startup()
1113+
1114+
try:
1115+
request = GenerationRequest(
1116+
content="test prompt",
1117+
request_type="text_completions",
1118+
constraints={"output_tokens": 50},
1119+
)
1120+
request_info = ScheduledRequestInfo(
1121+
request_id="test-id",
1122+
status="pending",
1123+
scheduler_node_id=1,
1124+
scheduler_process_id=1,
1125+
scheduler_start_time=123.0,
1126+
request_timings=GenerationRequestTimings(),
1127+
)
1128+
1129+
# Mock text_completions to test timing edge cases
1130+
async def mock_text_completions(*args, **kwargs):
1131+
yield None, None # Initial yield - tests line 343
1132+
yield "token1", None # First token
1133+
yield "token2", UsageStats(prompt_tokens=10, output_tokens=2) # Final
1134+
1135+
with patch.object(
1136+
backend, "text_completions", side_effect=mock_text_completions
1137+
):
1138+
responses = []
1139+
async for response, info in backend.resolve(request, request_info):
1140+
responses.append((response, info))
1141+
1142+
# Check that timing was properly set
1143+
final_response, final_info = responses[-1]
1144+
assert final_info.request_timings.request_start is not None
1145+
assert final_info.request_timings.first_iteration is not None
1146+
assert final_info.request_timings.last_iteration is not None
1147+
assert final_info.request_timings.request_end is not None
1148+
assert final_response.delta is None # Tests line 362
1149+
1150+
finally:
1151+
await backend.process_shutdown()

0 commit comments

Comments
 (0)