89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
from photools3.llm_session import LlmSession
|
|
|
|
|
|
def test_llm_session_mocked(test_photo):
|
|
# We want to mock OpenAI client's chat.completions.create
|
|
with patch("photools3.llm_session.OpenAI") as mocked_openai:
|
|
# Setup the mock response structure
|
|
mock_client = MagicMock()
|
|
mocked_openai.return_value = mock_client
|
|
|
|
# Mock the models list to return a proper object with an ID string
|
|
# This prevents Pydantic ValidationError when LlmSession tries to use model_id
|
|
mock_model = MagicMock()
|
|
mock_model.id = "gpt-4o"
|
|
mock_client.models.list.return_value.data = [mock_model]
|
|
|
|
# Create a response object: response.choices[0].message.content
|
|
mock_response = MagicMock()
|
|
mock_choice = MagicMock()
|
|
mock_message = MagicMock()
|
|
|
|
mock_message.content = (
|
|
'{"caption": "A beautiful sunset", "tags": ["sunset", "nature"]}'
|
|
)
|
|
mock_choice.message = mock_message
|
|
mock_response.choices = [mock_choice]
|
|
|
|
# Configure the client to return our mock response
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
# Create a new LlmSession instance that will use the mocked OpenAI client
|
|
mocked_session = LlmSession(system_prompt="You are a helper.")
|
|
|
|
# Run the test
|
|
result = mocked_session.call(
|
|
prompt="keywords for the attached picture", photos=[test_photo]
|
|
)
|
|
|
|
# Verify the results
|
|
assert result["caption"] == "A beautiful sunset"
|
|
assert "sunset" in result["tags"]
|
|
assert "nature" in result["tags"]
|
|
|
|
# Verify that the mock was called
|
|
assert mock_client.chat.completions.create.called
|
|
|
|
|
|
def test_llm_session_captions_and_keywords_mocked(test_photo):
|
|
# We want to mock OpenAI client's chat.completions.create
|
|
with patch("photools3.llm_session.OpenAI") as mocked_openai:
|
|
# Setup the mock response structure
|
|
mock_client = MagicMock()
|
|
mocked_openai.return_value = mock_client
|
|
|
|
# Mock the models list to return a proper object with an ID string
|
|
# This prevents Pydantic ValidationError when LlmSession tries to use model_id
|
|
mock_model = MagicMock()
|
|
mock_model.id = "gpt-4o"
|
|
mock_client.models.list.return_value.data = [mock_model]
|
|
|
|
# Create a response object: response.choices[0].message.content
|
|
mock_response = MagicMock()
|
|
mock_choice = MagicMock()
|
|
mock_message = MagicMock()
|
|
|
|
mock_message.content = '{"caption": "A mocked caption", "tags": ["mocked_tag"]}'
|
|
mock_choice.message = mock_message
|
|
mock_response.choices = [mock_choice]
|
|
|
|
# Configure the client to return our mock response
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
# Create a new LlmSessionCaptionsAndKeywords instance that will use the mocked OpenAI client
|
|
from photools3.llm_session import LlmSessionCaptionsAndKeywords
|
|
|
|
mocked_session = LlmSessionCaptionsAndKeywords()
|
|
|
|
# Run the test
|
|
result = mocked_session.call(photos=[test_photo])
|
|
|
|
# Verify the results
|
|
assert result["caption"] == "A mocked caption"
|
|
assert "mocked_tag" in result["tags"]
|
|
|
|
# Verify that the mock was called
|
|
assert mock_client.chat.completions.create.called
|