Fix tiktoken mock: use class not instance, remove TESTING check

This commit is contained in:
Russell Ballestrini 2026-01-11 09:17:49 -05:00
parent 9fb7047866
commit cf63e5c6cd
2 changed files with 7 additions and 6 deletions

View file

@ -125,7 +125,7 @@ class Message(db.Model):
if self.token_count is None:
if self.is_base64_image():
self.token_count = 0
elif not TIKTOKEN_AVAILABLE or os.environ.get("TESTING"):
elif not TIKTOKEN_AVAILABLE:
# Fallback: estimate ~4 chars per token when tiktoken unavailable
self.token_count = len(self.content) // 4 + 1
else:

View file

@ -36,8 +36,9 @@ class MockTiktoken:
# Insert mock tiktoken into sys.modules BEFORE any imports
# Use the class itself (not an instance) so patching works correctly
if 'tiktoken' not in sys.modules:
sys.modules['tiktoken'] = MockTiktoken()
sys.modules['tiktoken'] = MockTiktoken
def pytest_configure(config):
@ -46,12 +47,12 @@ def pytest_configure(config):
Ensures tiktoken is mocked before any test imports happen.
"""
if 'tiktoken' not in sys.modules:
sys.modules['tiktoken'] = MockTiktoken()
sys.modules['tiktoken'] = MockTiktoken
else:
# If tiktoken was already imported, patch its functions
import tiktoken
tiktoken.encoding_for_model = MockTiktoken.encoding_for_model
tiktoken.get_encoding = MockTiktoken.get_encoding
tiktoken_mod = sys.modules['tiktoken']
tiktoken_mod.encoding_for_model = MockTiktoken.encoding_for_model
tiktoken_mod.get_encoding = MockTiktoken.get_encoding
# =============================================================================