Voucher-Bot / tests /test_message_classification.py
Raj718's picture
Fix dark mode input font color and improve UI accessibility - Fixed dark mode input textbox font color for better visibility - Increased chatbot response font size - Enhanced Spanish language support for search intent detection - Improved email request detection and transit/school context handling - Integrated violation checker into search flow - Added comprehensive test suites for regex and LLM fallback validation
67e153c
import unittest
from email_handler import enhanced_classify_message
class TestMessageClassification(unittest.TestCase):
def setUp(self):
"""Set up test data before each test"""
self.empty_state = {}
self.state_with_listings = {
"listings": [
{
"address": "123 Test Ave, Bronx, NY 10457",
"price": "$2,000",
"url": "https://test.com/listing1",
"risk_level": "✅",
"building_violations": 0,
"title": "Nice Bronx Apartment"
},
{
"address": "456 Sample St, Brooklyn, NY 11201",
"price": "$2,500",
"url": "https://test.com/listing2",
"risk_level": "⚠️",
"building_violations": 2,
"title": "Cozy Brooklyn Spot"
}
]
}
def test_email_requests(self):
"""Test various forms of email requests"""
# Test email requests with listings
self.assertEqual(
enhanced_classify_message("can you write an email for listing #1", self.state_with_listings),
"email_request"
)
self.assertEqual(
enhanced_classify_message("write an email to the landlord of listing 1", self.state_with_listings),
"email_request"
)
self.assertEqual(
enhanced_classify_message("compose email for the first listing", self.state_with_listings),
"email_request"
)
self.assertEqual(
enhanced_classify_message("email the owner of listing #2", self.state_with_listings),
"email_request"
)
self.assertEqual(
enhanced_classify_message("contact listing 1", self.state_with_listings),
"email_request"
)
self.assertEqual(
enhanced_classify_message("send a message to listing 2", self.state_with_listings),
"email_request"
)
# Test contextual email requests (NEW)
self.assertEqual(
enhanced_classify_message("can you write an email for this one? My name is bob ross", self.state_with_listings),
"email_request"
)
self.assertEqual(
enhanced_classify_message("write an email for this listing", self.state_with_listings),
"email_request"
)
self.assertEqual(
enhanced_classify_message("compose email for this property", self.state_with_listings),
"email_request"
)
self.assertEqual(
enhanced_classify_message("email the owner of this one", self.state_with_listings),
"email_request"
)
self.assertEqual(
enhanced_classify_message("contact this listing", self.state_with_listings),
"email_request"
)
# Test that email requests without listing references don't work
self.assertEqual(
enhanced_classify_message("write to the landlord", self.state_with_listings),
"general_conversation"
)
self.assertEqual(
enhanced_classify_message("can you send an email", self.state_with_listings),
"general_conversation"
)
# Test email requests without listings in state
self.assertEqual(
enhanced_classify_message("can you write an email for listing #1", self.empty_state),
"general_conversation"
)
self.assertEqual(
enhanced_classify_message("can you write an email for this one", self.empty_state),
"general_conversation"
)
def test_search_requests(self):
"""Test various forms of search requests"""
# Test explicit search patterns
self.assertEqual(
enhanced_classify_message("find me an apartment", self.empty_state),
"new_search"
)
self.assertEqual(
enhanced_classify_message("search for listings", self.empty_state),
"new_search"
)
self.assertEqual(
enhanced_classify_message("i'm looking for a place", self.empty_state),
"new_search"
)
# Test location-based searches
self.assertEqual(
enhanced_classify_message("i have a section 8 voucher in the bronx", self.empty_state),
"new_search"
)
self.assertEqual(
enhanced_classify_message("looking for apartments in brooklyn", self.empty_state),
"new_search"
)
self.assertEqual(
enhanced_classify_message("need a 2 bedroom in manhattan", self.empty_state),
"new_search"
)
# Test that general location questions don't trigger search
self.assertEqual(
enhanced_classify_message("can i use my voucher in brooklyn?", self.empty_state),
"general_conversation"
)
self.assertEqual(
enhanced_classify_message("do they accept section 8 in manhattan?", self.empty_state),
"general_conversation"
)
def test_listing_requests(self):
"""Test various forms of listing requests"""
# Test numeric requests
self.assertEqual(
enhanced_classify_message("1", self.state_with_listings),
"listing_question"
)
self.assertEqual(
enhanced_classify_message("show me #2", self.state_with_listings),
"listing_question"
)
self.assertEqual(
enhanced_classify_message("can i see listing 1", self.state_with_listings),
"listing_question"
)
self.assertEqual(
enhanced_classify_message("what about listing #2", self.state_with_listings),
"listing_question"
)
# Test that listing requests require listings in state
self.assertEqual(
enhanced_classify_message("1", self.empty_state),
"general_conversation"
)
self.assertEqual(
enhanced_classify_message("show me listing 1", self.empty_state),
"general_conversation"
)
# Test invalid listing numbers
self.assertEqual(
enhanced_classify_message("show me listing 0", self.state_with_listings),
"listing_question" # Still returns listing_question, validation happens later
)
self.assertEqual(
enhanced_classify_message("can i see listing 999", self.state_with_listings),
"listing_question" # Still returns listing_question, validation happens later
)
def test_mixed_requests(self):
"""Test edge cases and mixed requests"""
# Test that search takes priority over listing when no listings exist
self.assertEqual(
enhanced_classify_message("show me listing 1 in brooklyn", self.empty_state),
"new_search"
)
# Test that listing request takes priority when listings exist
self.assertEqual(
enhanced_classify_message("show me listing 1 in brooklyn", self.state_with_listings),
"listing_question"
)
# Test that general questions don't trigger listing or search
self.assertEqual(
enhanced_classify_message("how does section 8 work?", self.state_with_listings),
"general_conversation"
)
self.assertEqual(
enhanced_classify_message("what documents do i need?", self.state_with_listings),
"general_conversation"
)
if __name__ == '__main__':
unittest.main()