mirror of
https://github.com/qodo-ai/pr-agent.git
synced 2025-12-12 10:55:17 +00:00
1c6958069a
- There is no test code for get_max_tokens of utils, so the unit test code related to the part corresponding to each test branch is added
67 lines
2 KiB
Python
67 lines
2 KiB
Python
import pytest
|
|
from pr_agent.algo.utils import get_max_tokens, MAX_TOKENS
|
|
import pr_agent.algo.utils as utils
|
|
|
|
class TestGetMaxTokens:
|
|
|
|
# Test if the file is in MAX_TOKENS
|
|
def test_model_max_tokens(self, monkeypatch):
|
|
fake_settings = type('', (), {
|
|
'config': type('', (), {
|
|
'custom_model_max_tokens': 0,
|
|
'max_model_tokens': 0
|
|
})()
|
|
})()
|
|
|
|
monkeypatch.setattr(utils, "get_settings", lambda: fake_settings)
|
|
|
|
model = "gpt-3.5-turbo"
|
|
expected = MAX_TOKENS[model]
|
|
|
|
assert get_max_tokens(model) == expected
|
|
|
|
# Test situations where the model is not registered and exists as a custom model
|
|
def test_model_has_custom(self, monkeypatch):
|
|
fake_settings = type('', (), {
|
|
'config': type('', (), {
|
|
'custom_model_max_tokens': 5000,
|
|
'max_model_tokens': 0 # 제한 없음
|
|
})()
|
|
})()
|
|
|
|
monkeypatch.setattr(utils, "get_settings", lambda: fake_settings)
|
|
|
|
model = "custom-model"
|
|
expected = 5000
|
|
|
|
assert get_max_tokens(model) == expected
|
|
|
|
def test_model_not_max_tokens_and_not_has_custom(self, monkeypatch):
|
|
fake_settings = type('', (), {
|
|
'config': type('', (), {
|
|
'custom_model_max_tokens': 0,
|
|
'max_model_tokens': 0
|
|
})()
|
|
})()
|
|
|
|
monkeypatch.setattr(utils, "get_settings", lambda: fake_settings)
|
|
|
|
model = "custom-model"
|
|
|
|
with pytest.raises(Exception):
|
|
get_max_tokens(model)
|
|
|
|
def test_model_max_tokens_with__limit(self, monkeypatch):
|
|
fake_settings = type('', (), {
|
|
'config': type('', (), {
|
|
'custom_model_max_tokens': 0,
|
|
'max_model_tokens': 10000
|
|
})()
|
|
})()
|
|
|
|
monkeypatch.setattr(utils, "get_settings", lambda: fake_settings)
|
|
|
|
model = "gpt-3.5-turbo" # this model setting is 160000
|
|
expected = 10000
|
|
|
|
assert get_max_tokens(model) == expected
|