"""Unit tests for company rating tool."""
import pytest
from unittest.mock import Mock, patch, AsyncMock
from fastmcp import Context
from birre.tools.company_rating import (
register_company_rating_tool,
get_company_rating_impl,
_calculate_letter_grade,
)
from birre.server import create_server
class TestCompanyRatingTool:
"""Test cases for company rating functionality."""
@pytest.fixture
def server(self):
"""Create a test server instance."""
return create_server()
@pytest.fixture
def mock_context(self):
"""Create a mock context for testing."""
context = Mock(spec=Context)
context.info = AsyncMock()
context.error = AsyncMock()
context.debug = AsyncMock()
context.warning = AsyncMock()
return context
@pytest.mark.asyncio
async def test_get_rating_already_subscribed_success(self, server, mock_context):
"""Test getting rating for company already subscribed."""
mock_company_details = {
"name": "Test Company",
"current_rating": 750,
"ratings": [{"rating_date": "2024-01-15"}],
}
with patch("birre.tools.company_rating.bitsight") as mock_bitsight:
mock_companies = Mock()
mock_companies.get_company_details.return_value = mock_company_details
mock_bitsight.Companies.return_value = mock_companies
result = await get_company_rating_impl(
mock_context, "test-guid-123")
assert result["rating"] == 750
assert result["rating_date"] == "2024-01-15"
assert result["grade"] == "A" # 750 should be grade A
assert result["company_name"] == "Test Company"
assert result["auto_subscribed"] is False
assert result["auto_unsubscribed"] is False
assert "error" not in result
# Verify company details was called
mock_companies.get_company_details.assert_called_once_with("test-guid-123")
@pytest.mark.asyncio
async def test_get_rating_auto_subscribe_success(self, server, mock_context):
"""Test auto-subscribe flow when company is not subscribed."""
# First call returns no rating (not subscribed)
mock_company_details_no_rating = {
"name": "Test Company",
"current_rating": None,
"ratings": [],
}
# Second call after subscription returns rating
mock_company_details_with_rating = {
"name": "Test Company",
"current_rating": 650,
"ratings": [{"rating_date": "2024-01-15"}],
}
with patch("birre.tools.company_rating.bitsight") as mock_bitsight:
mock_companies = Mock()
mock_subscriptions = Mock()
# First call has no rating, second call has rating
mock_companies.get_company_details.side_effect = [
mock_company_details_no_rating,
mock_company_details_with_rating,
]
mock_subscriptions.post_subscribe.return_value = {"status": "success"}
mock_bitsight.Companies.return_value = mock_companies
mock_bitsight.Subscriptions.return_value = mock_subscriptions
result = await get_company_rating_impl(
mock_context, "test-guid-123")
assert result["rating"] == 650
assert result["grade"] == "B" # 650 should be grade B
assert result["company_name"] == "Test Company"
assert result["auto_subscribed"] is True
assert (
result["auto_unsubscribed"] is True
) # Should attempt unsubscribe
# Verify subscription was called
mock_subscriptions.post_subscribe.assert_called_once_with(
"test-guid-123"
)
@pytest.mark.asyncio
async def test_get_rating_invalid_guid(self, server, mock_context):
"""Test error handling for invalid GUID."""
# Test empty GUID
result = await get_company_rating_impl(mock_context, "")
assert "error" in result
assert "Valid company GUID is required" in result["error"]
assert result["rating"] is None
# Test None GUID
result = await get_company_rating_impl(mock_context, None)
assert "error" in result
assert "Valid company GUID is required" in result["error"]
@pytest.mark.asyncio
async def test_get_rating_subscription_failure(self, server, mock_context):
"""Test handling when subscription fails."""
with patch("birre.tools.company_rating.bitsight") as mock_bitsight:
mock_companies = Mock()
mock_subscriptions = Mock()
# First call raises exception (needs subscription)
mock_companies.get_company_details.side_effect = [
Exception("Not subscribed"),
# Won't get to second call due to subscription failure
]
# Subscription fails
mock_subscriptions.post_subscribe.side_effect = Exception(
"Subscription failed"
)
mock_bitsight.Companies.return_value = mock_companies
mock_bitsight.Subscriptions.return_value = mock_subscriptions
result = await get_company_rating_impl(
mock_context, "test-guid-123")
assert "error" in result
assert "Failed to subscribe to company" in result["error"]
assert result["rating"] is None
assert result["auto_subscribed"] is False
assert result["auto_unsubscribed"] is False
@pytest.mark.asyncio
async def test_get_rating_no_rating_after_subscribe(self, server, mock_context):
"""Test when rating is not immediately available after subscription."""
mock_company_details_no_rating = {
"name": "Test Company",
"rating": None,
"rating_date": None,
}
with patch("birre.tools.company_rating.bitsight") as mock_bitsight:
mock_companies = Mock()
mock_subscriptions = Mock()
# Both calls return no rating
mock_companies.get_company_details.side_effect = [
Exception("Not subscribed"),
mock_company_details_no_rating, # Still no rating after subscription
]
mock_subscriptions.post_subscribe.return_value = {"status": "success"}
mock_bitsight.Companies.return_value = mock_companies
mock_bitsight.Subscriptions.return_value = mock_subscriptions
result = await get_company_rating_impl(
mock_context, "test-guid-123")
assert "error" in result
assert (
"Rating data not available immediately after subscription"
in result["error"]
)
assert result["rating"] is None
assert result["auto_subscribed"] is True
assert (
result["auto_unsubscribed"] is False
) # Won't unsubscribe if no rating
@pytest.mark.asyncio
async def test_get_rating_general_exception(self, server, mock_context):
"""Test handling of unexpected exceptions."""
with patch("birre.tools.company_rating.bitsight") as mock_bitsight:
mock_bitsight.Companies.side_effect = Exception("Unexpected error")
result = await get_company_rating_impl(
mock_context, "test-guid-123")
assert "error" in result
assert (
"Failed to get company rating: Unexpected error" in result["error"]
)
assert result["auto_subscribed"] is False
assert result["auto_unsubscribed"] is False
class TestLetterGradeCalculation:
"""Test cases for letter grade calculation."""
def test_calculate_letter_grade_a(self):
"""Test A grade calculation."""
assert _calculate_letter_grade(900) == "A"
assert _calculate_letter_grade(740) == "A" # Boundary
assert _calculate_letter_grade(800) == "A"
def test_calculate_letter_grade_b(self):
"""Test B grade calculation."""
assert _calculate_letter_grade(739) == "B"
assert _calculate_letter_grade(640) == "B" # Boundary
assert _calculate_letter_grade(700) == "B"
def test_calculate_letter_grade_c(self):
"""Test C grade calculation."""
assert _calculate_letter_grade(639) == "C"
assert _calculate_letter_grade(560) == "C" # Boundary
assert _calculate_letter_grade(600) == "C"
def test_calculate_letter_grade_d(self):
"""Test D grade calculation."""
assert _calculate_letter_grade(559) == "D"
assert _calculate_letter_grade(480) == "D" # Boundary
assert _calculate_letter_grade(520) == "D"
def test_calculate_letter_grade_f(self):
"""Test F grade calculation."""
assert _calculate_letter_grade(479) == "F"
assert _calculate_letter_grade(0) == "F"
assert _calculate_letter_grade(300) == "F"
