style: Fix linting issues in test_models.py

This commit is contained in:
Paul Gauthier (aider) 2025-03-11 11:30:28 -07:00
parent 58cd190ca9
commit 68c27f885f

View file

@ -159,48 +159,48 @@ class TestModels(unittest.TestCase):
model = Model("github/o1-preview") model = Model("github/o1-preview")
self.assertEqual(model.name, "github/o1-preview") self.assertEqual(model.name, "github/o1-preview")
self.assertEqual(model.use_temperature, False) self.assertEqual(model.use_temperature, False)
def test_parse_token_value(self): def test_parse_token_value(self):
# Create a model instance to test the parse_token_value method # Create a model instance to test the parse_token_value method
model = Model("gpt-4") model = Model("gpt-4")
# Test integer inputs # Test integer inputs
self.assertEqual(model.parse_token_value(8096), 8096) self.assertEqual(model.parse_token_value(8096), 8096)
self.assertEqual(model.parse_token_value(1000), 1000) self.assertEqual(model.parse_token_value(1000), 1000)
# Test string inputs # Test string inputs
self.assertEqual(model.parse_token_value("8096"), 8096) self.assertEqual(model.parse_token_value("8096"), 8096)
# Test k/K suffix (kilobytes) # Test k/K suffix (kilobytes)
self.assertEqual(model.parse_token_value("8k"), 8 * 1024) self.assertEqual(model.parse_token_value("8k"), 8 * 1024)
self.assertEqual(model.parse_token_value("8K"), 8 * 1024) self.assertEqual(model.parse_token_value("8K"), 8 * 1024)
self.assertEqual(model.parse_token_value("10.5k"), 10.5 * 1024) self.assertEqual(model.parse_token_value("10.5k"), 10.5 * 1024)
self.assertEqual(model.parse_token_value("0.5K"), 0.5 * 1024) self.assertEqual(model.parse_token_value("0.5K"), 0.5 * 1024)
# Test m/M suffix (megabytes) # Test m/M suffix (megabytes)
self.assertEqual(model.parse_token_value("1m"), 1 * 1024 * 1024) self.assertEqual(model.parse_token_value("1m"), 1 * 1024 * 1024)
self.assertEqual(model.parse_token_value("1M"), 1 * 1024 * 1024) self.assertEqual(model.parse_token_value("1M"), 1 * 1024 * 1024)
self.assertEqual(model.parse_token_value("0.5M"), 0.5 * 1024 * 1024) self.assertEqual(model.parse_token_value("0.5M"), 0.5 * 1024 * 1024)
# Test with spaces # Test with spaces
self.assertEqual(model.parse_token_value(" 8k "), 8 * 1024) self.assertEqual(model.parse_token_value(" 8k "), 8 * 1024)
# Test conversion from other types # Test conversion from other types
self.assertEqual(model.parse_token_value(8.0), 8) self.assertEqual(model.parse_token_value(8.0), 8)
def test_set_thinking_tokens(self): def test_set_thinking_tokens(self):
# Test that set_thinking_tokens correctly sets the tokens with different formats # Test that set_thinking_tokens correctly sets the tokens with different formats
model = Model("gpt-4") model = Model("gpt-4")
# Test with integer # Test with integer
model.set_thinking_tokens(8096) model.set_thinking_tokens(8096)
self.assertEqual(model.extra_params["thinking"]["budget_tokens"], 8096) self.assertEqual(model.extra_params["thinking"]["budget_tokens"], 8096)
self.assertFalse(model.use_temperature) self.assertFalse(model.use_temperature)
# Test with string # Test with string
model.set_thinking_tokens("10k") model.set_thinking_tokens("10k")
self.assertEqual(model.extra_params["thinking"]["budget_tokens"], 10 * 1024) self.assertEqual(model.extra_params["thinking"]["budget_tokens"], 10 * 1024)
# Test with decimal value # Test with decimal value
model.set_thinking_tokens("0.5M") model.set_thinking_tokens("0.5M")
self.assertEqual(model.extra_params["thinking"]["budget_tokens"], 0.5 * 1024 * 1024) self.assertEqual(model.extra_params["thinking"]["budget_tokens"], 0.5 * 1024 * 1024)