Spaces:
Sleeping
Sleeping
File size: 7,767 Bytes
67e153c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 |
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() |