michael
2ee0d1638a
- Replace ChatPromptTemplate with direct HumanMessage/SystemMessage to avoid template variable parsing - Fix f-string formatting issues in prompt strings - Update test_general_summary_fallback_on_error to properly mock fallback chain - Add tests directory with comprehensive test coverage
418 lines
17 KiB
Python
418 lines
17 KiB
Python
"""
|
|
Unit tests for generate_summary.py module
|
|
"""
|
|
import pytest
|
|
import os
|
|
import json
|
|
from unittest.mock import Mock, MagicMock, patch, call
|
|
from pydantic import ValidationError
|
|
|
|
# Import the functions and models to test
|
|
import sys
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
|
|
|
|
from scripts.generate_summary import (
|
|
general_summary,
|
|
custom_summary,
|
|
BasicSummary,
|
|
AdvancedSummary,
|
|
KeyPoint,
|
|
Summary,
|
|
Purpose,
|
|
Chapters,
|
|
Outcomes,
|
|
ActionItemsPerUser
|
|
)
|
|
|
|
|
|
class TestPydanticModels:
|
|
"""Test Pydantic model validation"""
|
|
|
|
def test_key_point_model(self):
|
|
"""Test KeyPoint model validation"""
|
|
key_point = KeyPoint(text="Test point", timestamp=10.5)
|
|
assert key_point.text == "Test point"
|
|
assert key_point.timestamp == 10.5
|
|
|
|
def test_summary_model(self):
|
|
"""Test Summary model validation"""
|
|
summary = Summary(text="Test summary", duration_minutes=15.5)
|
|
assert summary.text == "Test summary"
|
|
assert summary.duration_minutes == 15.5
|
|
|
|
def test_basic_summary_model(self, sample_basic_summary_response):
|
|
"""Test BasicSummary model validation"""
|
|
basic_summary = BasicSummary(**sample_basic_summary_response)
|
|
assert len(basic_summary.Key_Points) == 2
|
|
assert basic_summary.Summary.text is not None
|
|
assert basic_summary.Summary.duration_minutes > 0
|
|
|
|
def test_advanced_summary_model(self, sample_advanced_summary_response):
|
|
"""Test AdvancedSummary model validation"""
|
|
advanced_summary = AdvancedSummary(**sample_advanced_summary_response)
|
|
assert advanced_summary.Purpose.text is not None
|
|
assert len(advanced_summary.Chapters.content) > 0
|
|
assert len(advanced_summary.Outcomes.content) > 0
|
|
assert len(advanced_summary.Action_Items_Per_User) > 0
|
|
|
|
|
|
class TestGeneralSummary:
|
|
"""Test general_summary function"""
|
|
|
|
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
|
@patch('scripts.generate_summary.ChatAnthropic')
|
|
@patch('scripts.generate_summary.ChatPromptTemplate')
|
|
def test_general_summary_freemium_plan(
|
|
self,
|
|
mock_prompt_template,
|
|
mock_chat_anthropic,
|
|
sample_transcription_dict,
|
|
sample_basic_summary_response
|
|
):
|
|
"""Test general_summary with freemium plan"""
|
|
# Setup mocks
|
|
mock_model = MagicMock()
|
|
mock_chain = MagicMock()
|
|
mock_structured_model = MagicMock()
|
|
|
|
# Create a mock BasicSummary instance
|
|
from scripts.generate_summary import BasicSummary
|
|
mock_result = BasicSummary(**sample_basic_summary_response)
|
|
|
|
mock_structured_model.invoke.return_value = mock_result
|
|
mock_model.with_structured_output.return_value = mock_structured_model
|
|
mock_prompt_template.from_messages.return_value = MagicMock()
|
|
mock_chat_anthropic.return_value = mock_model
|
|
|
|
# Mock the pipe operator
|
|
mock_prompt_instance = MagicMock()
|
|
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
|
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
|
mock_chain.invoke.return_value = mock_result
|
|
|
|
# Call the function
|
|
result = general_summary(sample_transcription_dict, plan_tier="freemium")
|
|
|
|
# Verify results
|
|
assert result is not None
|
|
assert "Key_Points" in result
|
|
assert "Summary" in result
|
|
assert len(result["Key_Points"]) > 0
|
|
|
|
# Verify model was initialized with correct parameters
|
|
mock_chat_anthropic.assert_called_once()
|
|
call_args = mock_chat_anthropic.call_args
|
|
assert call_args.kwargs["model"] == "claude-sonnet-4-5-20250929"
|
|
assert call_args.kwargs["max_tokens"] == 2000
|
|
assert call_args.kwargs["temperature"] == 0.2
|
|
|
|
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
|
@patch('scripts.generate_summary.ChatAnthropic')
|
|
@patch('scripts.generate_summary.ChatPromptTemplate')
|
|
def test_general_summary_pro_plan(
|
|
self,
|
|
mock_prompt_template,
|
|
mock_chat_anthropic,
|
|
sample_transcription_dict,
|
|
sample_advanced_summary_response
|
|
):
|
|
"""Test general_summary with pro plan"""
|
|
# Setup mocks
|
|
mock_model = MagicMock()
|
|
mock_chain = MagicMock()
|
|
|
|
# Create a mock AdvancedSummary instance
|
|
from scripts.generate_summary import AdvancedSummary
|
|
mock_result = AdvancedSummary(**sample_advanced_summary_response)
|
|
|
|
mock_structured_model = MagicMock()
|
|
mock_structured_model.invoke.return_value = mock_result
|
|
mock_model.with_structured_output.return_value = mock_structured_model
|
|
mock_prompt_template.from_messages.return_value = MagicMock()
|
|
mock_chat_anthropic.return_value = mock_model
|
|
|
|
# Mock the pipe operator
|
|
mock_prompt_instance = MagicMock()
|
|
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
|
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
|
mock_chain.invoke.return_value = mock_result
|
|
|
|
# Call the function
|
|
result = general_summary(sample_transcription_dict, plan_tier="pro")
|
|
|
|
# Verify results
|
|
assert result is not None
|
|
assert "Purpose" in result
|
|
assert "Chapters" in result
|
|
assert "Outcomes" in result
|
|
assert "Action_Items_Per_User" in result
|
|
|
|
# Verify model was initialized with correct parameters
|
|
mock_chat_anthropic.assert_called_once()
|
|
call_args = mock_chat_anthropic.call_args
|
|
assert call_args.kwargs["model"] == "claude-sonnet-4-5-20250929"
|
|
assert call_args.kwargs["max_tokens"] == 4000
|
|
|
|
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
|
@patch('scripts.generate_summary.ChatAnthropic')
|
|
@patch('scripts.generate_summary.ChatPromptTemplate')
|
|
def test_general_summary_with_string_transcription(
|
|
self,
|
|
mock_prompt_template,
|
|
mock_chat_anthropic,
|
|
sample_basic_summary_response
|
|
):
|
|
"""Test general_summary with string transcription"""
|
|
transcription_str = json.dumps({"sentences": []})
|
|
|
|
mock_model = MagicMock()
|
|
mock_chain = MagicMock()
|
|
from scripts.generate_summary import BasicSummary
|
|
mock_result = BasicSummary(**sample_basic_summary_response)
|
|
|
|
mock_structured_model = MagicMock()
|
|
mock_structured_model.invoke.return_value = mock_result
|
|
mock_model.with_structured_output.return_value = mock_structured_model
|
|
mock_prompt_template.from_messages.return_value = MagicMock()
|
|
mock_chat_anthropic.return_value = mock_model
|
|
|
|
mock_prompt_instance = MagicMock()
|
|
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
|
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
|
mock_chain.invoke.return_value = mock_result
|
|
|
|
result = general_summary(transcription_str, plan_tier="freemium")
|
|
|
|
assert result is not None
|
|
assert "Key_Points" in result
|
|
|
|
def test_general_summary_missing_api_key(self, sample_transcription_dict):
|
|
"""Test general_summary raises error when API key is missing"""
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"):
|
|
general_summary(sample_transcription_dict, plan_tier="freemium")
|
|
|
|
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
|
@patch('scripts.generate_summary.ChatAnthropic')
|
|
@patch('scripts.generate_summary.ChatPromptTemplate')
|
|
def test_general_summary_fallback_on_error(
|
|
self,
|
|
mock_prompt_template,
|
|
mock_chat_anthropic,
|
|
sample_transcription_dict,
|
|
sample_basic_summary_response
|
|
):
|
|
"""Test general_summary falls back to non-structured output on error"""
|
|
mock_model = MagicMock()
|
|
mock_chain = MagicMock()
|
|
mock_fallback_chain = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.content = json.dumps(sample_basic_summary_response)
|
|
|
|
# First call (structured) raises error
|
|
mock_structured_model = MagicMock()
|
|
mock_structured_model.invoke.side_effect = Exception("Structured output failed")
|
|
mock_model.with_structured_output.return_value = mock_structured_model
|
|
|
|
# Set up the chain so that structured chain raises exception
|
|
# and fallback chain returns the mock_response
|
|
mock_prompt_instance = MagicMock()
|
|
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
|
|
|
# When __or__ is called with structured_model, return chain that raises exception
|
|
# When __or__ is called with model (fallback), return fallback_chain that succeeds
|
|
def or_handler(self, other):
|
|
if other == mock_structured_model:
|
|
# Structured chain - should raise exception
|
|
mock_chain.invoke.side_effect = Exception("Structured output failed")
|
|
return mock_chain
|
|
elif other == mock_model:
|
|
# Fallback chain - should succeed
|
|
mock_fallback_chain.invoke.return_value = mock_response
|
|
return mock_fallback_chain
|
|
return mock_chain
|
|
|
|
mock_prompt_instance.__or__ = or_handler
|
|
mock_chat_anthropic.return_value = mock_model
|
|
|
|
result = general_summary(sample_transcription_dict, plan_tier="freemium")
|
|
|
|
assert result is not None
|
|
assert "Key_Points" in result
|
|
|
|
|
|
class TestCustomSummary:
|
|
"""Test custom_summary function"""
|
|
|
|
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
|
@patch('scripts.generate_summary.ChatAnthropic')
|
|
@patch('scripts.generate_summary.ChatPromptTemplate')
|
|
def test_custom_summary_success(
|
|
self,
|
|
mock_prompt_template,
|
|
mock_chat_anthropic,
|
|
sample_transcription_dict,
|
|
sample_template
|
|
):
|
|
"""Test custom_summary with successful response"""
|
|
mock_model = MagicMock()
|
|
mock_chain = MagicMock()
|
|
mock_response = MagicMock()
|
|
|
|
expected_result = {
|
|
"Key_Points": {"content": []},
|
|
"Summary": {"content": []},
|
|
"Next_Steps": {"content": []}
|
|
}
|
|
mock_response.content = json.dumps(expected_result)
|
|
|
|
mock_prompt_instance = MagicMock()
|
|
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
|
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
|
mock_chain.invoke.return_value = mock_response
|
|
mock_chat_anthropic.return_value = mock_model
|
|
|
|
result = custom_summary(sample_template, sample_transcription_dict)
|
|
|
|
assert result is not None
|
|
assert "Key_Points" in result or "Summary" in result
|
|
|
|
# Verify model was initialized correctly
|
|
mock_chat_anthropic.assert_called_once()
|
|
call_args = mock_chat_anthropic.call_args
|
|
assert call_args.kwargs["model"] == "claude-sonnet-4-5-20250929"
|
|
assert call_args.kwargs["max_tokens"] == 8000
|
|
|
|
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
|
@patch('scripts.generate_summary.ChatAnthropic')
|
|
@patch('scripts.generate_summary.ChatPromptTemplate')
|
|
def test_custom_summary_with_markdown_wrapper(
|
|
self,
|
|
mock_prompt_template,
|
|
mock_chat_anthropic,
|
|
sample_transcription_dict,
|
|
sample_template
|
|
):
|
|
"""Test custom_summary handles markdown-wrapped JSON"""
|
|
mock_model = MagicMock()
|
|
mock_chain = MagicMock()
|
|
mock_response = MagicMock()
|
|
|
|
expected_result = {"result": "test"}
|
|
wrapped_json = f"```json\n{json.dumps(expected_result)}\n```"
|
|
mock_response.content = wrapped_json
|
|
|
|
mock_prompt_instance = MagicMock()
|
|
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
|
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
|
mock_chain.invoke.return_value = mock_response
|
|
mock_chat_anthropic.return_value = mock_model
|
|
|
|
result = custom_summary(sample_template, sample_transcription_dict)
|
|
|
|
assert result == expected_result
|
|
|
|
def test_custom_summary_missing_api_key(self, sample_transcription_dict, sample_template):
|
|
"""Test custom_summary raises error when API key is missing"""
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"):
|
|
custom_summary(sample_template, sample_transcription_dict)
|
|
|
|
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
|
@patch('scripts.generate_summary.ChatAnthropic')
|
|
@patch('scripts.generate_summary.ChatPromptTemplate')
|
|
def test_custom_summary_invalid_json(
|
|
self,
|
|
mock_prompt_template,
|
|
mock_chat_anthropic,
|
|
sample_transcription_dict,
|
|
sample_template
|
|
):
|
|
"""Test custom_summary handles invalid JSON gracefully"""
|
|
mock_model = MagicMock()
|
|
mock_chain = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.content = "This is not valid JSON"
|
|
|
|
mock_prompt_instance = MagicMock()
|
|
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
|
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
|
mock_chain.invoke.return_value = mock_response
|
|
mock_chat_anthropic.return_value = mock_model
|
|
|
|
with pytest.raises(ValueError, match="Could not parse response as JSON"):
|
|
custom_summary(sample_template, sample_transcription_dict)
|
|
|
|
|
|
class TestSchemaSwitching:
|
|
"""Test schema switching based on plan tier"""
|
|
|
|
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
|
@patch('scripts.generate_summary.ChatAnthropic')
|
|
@patch('scripts.generate_summary.ChatPromptTemplate')
|
|
def test_schema_switching_freemium(
|
|
self,
|
|
mock_prompt_template,
|
|
mock_chat_anthropic,
|
|
sample_transcription_dict,
|
|
sample_basic_summary_response
|
|
):
|
|
"""Test that freemium plan uses BasicSummary schema"""
|
|
mock_model = MagicMock()
|
|
mock_chain = MagicMock()
|
|
from scripts.generate_summary import BasicSummary
|
|
mock_result = BasicSummary(**sample_basic_summary_response)
|
|
|
|
mock_structured_model = MagicMock()
|
|
mock_structured_model.invoke.return_value = mock_result
|
|
mock_model.with_structured_output = MagicMock(return_value=mock_structured_model)
|
|
|
|
mock_prompt_instance = MagicMock()
|
|
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
|
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
|
mock_chain.invoke.return_value = mock_result
|
|
mock_chat_anthropic.return_value = mock_model
|
|
|
|
general_summary(sample_transcription_dict, plan_tier="freemium")
|
|
|
|
# Verify BasicSummary schema was used
|
|
mock_model.with_structured_output.assert_called_once()
|
|
call_args = mock_model.with_structured_output.call_args
|
|
from scripts.generate_summary import BasicSummary
|
|
assert call_args[0][0] == BasicSummary
|
|
|
|
@patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'test-key'})
|
|
@patch('scripts.generate_summary.ChatAnthropic')
|
|
@patch('scripts.generate_summary.ChatPromptTemplate')
|
|
def test_schema_switching_pro(
|
|
self,
|
|
mock_prompt_template,
|
|
mock_chat_anthropic,
|
|
sample_transcription_dict,
|
|
sample_advanced_summary_response
|
|
):
|
|
"""Test that pro plan uses AdvancedSummary schema"""
|
|
mock_model = MagicMock()
|
|
mock_chain = MagicMock()
|
|
from scripts.generate_summary import AdvancedSummary
|
|
mock_result = AdvancedSummary(**sample_advanced_summary_response)
|
|
|
|
mock_structured_model = MagicMock()
|
|
mock_structured_model.invoke.return_value = mock_result
|
|
mock_model.with_structured_output = MagicMock(return_value=mock_structured_model)
|
|
|
|
mock_prompt_instance = MagicMock()
|
|
mock_prompt_template.from_messages.return_value = mock_prompt_instance
|
|
mock_prompt_instance.__or__ = lambda self, other: mock_chain
|
|
mock_chain.invoke.return_value = mock_result
|
|
mock_chat_anthropic.return_value = mock_model
|
|
|
|
general_summary(sample_transcription_dict, plan_tier="pro")
|
|
|
|
# Verify AdvancedSummary schema was used
|
|
mock_model.with_structured_output.assert_called_once()
|
|
call_args = mock_model.with_structured_output.call_args
|
|
from scripts.generate_summary import AdvancedSummary
|
|
assert call_args[0][0] == AdvancedSummary
|
|
|