refactor: isolate database migrations and model defaults
This commit is contained in:
@@ -19,7 +19,8 @@ class AuthenticationGuardTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self) -> None:
|
||||
self.temporary_directory = tempfile.TemporaryDirectory()
|
||||
self.storage = AssistantStorage(
|
||||
Path(self.temporary_directory.name) / "assistant.sqlite3"
|
||||
Path(self.temporary_directory.name) / "assistant.sqlite3",
|
||||
default_models={"local": "qwen3.5:9b"},
|
||||
)
|
||||
|
||||
def tearDown(self) -> None:
|
||||
|
||||
@@ -6,6 +6,20 @@ from assistant_bot import config
|
||||
|
||||
|
||||
class AppSettingsTests(unittest.TestCase):
|
||||
@staticmethod
|
||||
def load_settings(**overrides):
|
||||
environment = {
|
||||
config.TOKEN_ENV_NAME: "token",
|
||||
config.PASSWORD_ENV_NAME: "secret",
|
||||
**overrides,
|
||||
}
|
||||
with patch("assistant_bot.config.load_env_file"), patch.dict(
|
||||
os.environ,
|
||||
environment,
|
||||
clear=True,
|
||||
):
|
||||
return config.load_app_settings()
|
||||
|
||||
def test_loads_complete_startup_configuration_once(self) -> None:
|
||||
environment = {
|
||||
config.TOKEN_ENV_NAME: "token",
|
||||
@@ -33,87 +47,59 @@ class AppSettingsTests(unittest.TestCase):
|
||||
self.assertIsNone(settings.whisper_language)
|
||||
self.assertIsNone(settings.yandex_cloud_folder)
|
||||
|
||||
|
||||
class WhisperConfigTests(unittest.TestCase):
|
||||
def test_gpu_int8_defaults(self) -> None:
|
||||
variable_names = (
|
||||
config.WHISPER_MODEL_ENV_NAME,
|
||||
config.WHISPER_DEVICE_ENV_NAME,
|
||||
config.WHISPER_COMPUTE_TYPE_ENV_NAME,
|
||||
)
|
||||
with patch("assistant_bot.config.load_env_file"), patch.dict(
|
||||
os.environ,
|
||||
{},
|
||||
clear=False,
|
||||
):
|
||||
for variable_name in variable_names:
|
||||
os.environ.pop(variable_name, None)
|
||||
settings = self.load_settings()
|
||||
|
||||
self.assertEqual(config.get_whisper_model(), "large-v3")
|
||||
self.assertEqual(config.get_whisper_device(), "cuda")
|
||||
self.assertEqual(config.get_whisper_compute_type(), "int8")
|
||||
self.assertEqual(settings.whisper_model, "large-v3")
|
||||
self.assertEqual(settings.whisper_device, "cuda")
|
||||
self.assertEqual(settings.whisper_compute_type, "int8")
|
||||
|
||||
|
||||
class AssistantModeConfigTests(unittest.TestCase):
|
||||
def test_local_mode_is_default(self) -> None:
|
||||
with patch("assistant_bot.config.load_env_file"), patch.dict(
|
||||
os.environ,
|
||||
{},
|
||||
clear=False,
|
||||
):
|
||||
os.environ.pop(config.MODE_ENV_NAME, None)
|
||||
self.assertEqual(config.get_assistant_mode(), "local")
|
||||
self.assertEqual(self.load_settings().mode, "local")
|
||||
|
||||
def test_yandex_mode_is_supported(self) -> None:
|
||||
with patch("assistant_bot.config.load_env_file"), patch.dict(
|
||||
os.environ,
|
||||
{config.MODE_ENV_NAME: "YANDEX"},
|
||||
clear=False,
|
||||
):
|
||||
self.assertEqual(config.get_assistant_mode(), "yandex")
|
||||
settings = self.load_settings(
|
||||
**{
|
||||
config.MODE_ENV_NAME: "YANDEX",
|
||||
config.YANDEX_CLOUD_FOLDER_ENV_NAME: "folder-id",
|
||||
}
|
||||
)
|
||||
|
||||
self.assertEqual(settings.mode, "yandex")
|
||||
|
||||
def test_unknown_mode_is_rejected(self) -> None:
|
||||
with patch("assistant_bot.config.load_env_file"), patch.dict(
|
||||
os.environ,
|
||||
{config.MODE_ENV_NAME: "cloud"},
|
||||
clear=False,
|
||||
):
|
||||
with self.assertRaises(RuntimeError):
|
||||
config.get_assistant_mode()
|
||||
with self.assertRaises(RuntimeError):
|
||||
self.load_settings(**{config.MODE_ENV_NAME: "cloud"})
|
||||
|
||||
def test_yandex_model_is_full_gpt_uri(self) -> None:
|
||||
with patch("assistant_bot.config.load_env_file"), patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
settings = self.load_settings(
|
||||
**{
|
||||
config.MODE_ENV_NAME: "yandex",
|
||||
config.YANDEX_CLOUD_FOLDER_ENV_NAME: "folder-id",
|
||||
config.YANDEX_CLOUD_MODEL_ENV_NAME: "yandexgpt/latest",
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
self.assertEqual(
|
||||
config.get_default_yandex_model(),
|
||||
"gpt://folder-id/yandexgpt/latest",
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
settings.default_models()["yandex"],
|
||||
"gpt://folder-id/yandexgpt/latest",
|
||||
)
|
||||
|
||||
class PasswordConfigTests(unittest.TestCase):
|
||||
def test_password_is_required(self) -> None:
|
||||
with patch("assistant_bot.config.load_env_file"), patch.dict(
|
||||
os.environ,
|
||||
{},
|
||||
clear=False,
|
||||
{config.TOKEN_ENV_NAME: "token"},
|
||||
clear=True,
|
||||
):
|
||||
os.environ.pop(config.PASSWORD_ENV_NAME, None)
|
||||
with self.assertRaises(RuntimeError):
|
||||
config.get_assistant_password()
|
||||
config.load_app_settings()
|
||||
|
||||
def test_password_is_read_from_environment(self) -> None:
|
||||
with patch("assistant_bot.config.load_env_file"), patch.dict(
|
||||
os.environ,
|
||||
{config.PASSWORD_ENV_NAME: " test-password "},
|
||||
clear=False,
|
||||
):
|
||||
self.assertEqual(config.get_assistant_password(), "test-password")
|
||||
settings = self.load_settings(
|
||||
**{config.PASSWORD_ENV_NAME: " test-password "}
|
||||
)
|
||||
|
||||
self.assertEqual(settings.assistant_password, "test-password")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -12,10 +12,21 @@ from assistant_bot.agent import (
|
||||
parse_agent_decision,
|
||||
)
|
||||
from assistant_bot.agent_tools import AGENT_TOOLS
|
||||
from assistant_bot.migrations import LATEST_SCHEMA_VERSION
|
||||
from assistant_bot.reminders import parse_reminder, unit_to_timedelta
|
||||
from assistant_bot.storage import AssistantStorage
|
||||
|
||||
|
||||
DEFAULT_MODELS = {
|
||||
"local": "qwen3.5:9b",
|
||||
"yandex": "gpt://folder/yandexgpt/latest",
|
||||
}
|
||||
|
||||
|
||||
def create_storage(path: Path) -> AssistantStorage:
|
||||
return AssistantStorage(path, default_models=DEFAULT_MODELS)
|
||||
|
||||
|
||||
class ReminderParserTests(unittest.TestCase):
|
||||
def test_supported_relative_unit(self) -> None:
|
||||
delta = unit_to_timedelta(2, "часа")
|
||||
@@ -89,7 +100,7 @@ class AgentToolContractTests(unittest.TestCase):
|
||||
|
||||
def test_crud_tool_result_shapes_remain_stable(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
storage = AssistantStorage(Path(directory) / "assistant.sqlite3")
|
||||
storage = create_storage(Path(directory) / "assistant.sqlite3")
|
||||
common = {
|
||||
"storage": storage,
|
||||
"user_id": 42,
|
||||
@@ -187,7 +198,7 @@ class AgentToolContractTests(unittest.TestCase):
|
||||
result = execute_agent_tool(
|
||||
"missing",
|
||||
{},
|
||||
AssistantStorage(Path(directory) / "assistant.sqlite3"),
|
||||
create_storage(Path(directory) / "assistant.sqlite3"),
|
||||
user_id=42,
|
||||
chat_id=100,
|
||||
tz=ZoneInfo("UTC"),
|
||||
@@ -200,15 +211,27 @@ class AgentToolContractTests(unittest.TestCase):
|
||||
|
||||
|
||||
class StorageTests(unittest.TestCase):
|
||||
def test_new_database_uses_latest_schema_version(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
database_path = Path(directory) / "assistant.sqlite3"
|
||||
create_storage(database_path)
|
||||
|
||||
with closing(sqlite3.connect(database_path)) as connection:
|
||||
version = connection.execute(
|
||||
"PRAGMA user_version"
|
||||
).fetchone()[0]
|
||||
|
||||
self.assertEqual(version, LATEST_SCHEMA_VERSION)
|
||||
|
||||
def test_user_authorization_is_persisted(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
database_path = Path(directory) / "assistant.sqlite3"
|
||||
storage = AssistantStorage(database_path)
|
||||
storage = create_storage(database_path)
|
||||
|
||||
self.assertFalse(storage.is_user_authorized(42))
|
||||
storage.authorize_user(42)
|
||||
|
||||
reopened_storage = AssistantStorage(database_path)
|
||||
reopened_storage = create_storage(database_path)
|
||||
self.assertTrue(reopened_storage.is_user_authorized(42))
|
||||
self.assertFalse(reopened_storage.is_user_authorized(43))
|
||||
|
||||
@@ -227,14 +250,32 @@ class StorageTests(unittest.TestCase):
|
||||
)
|
||||
connection.commit()
|
||||
|
||||
storage = AssistantStorage(database_path)
|
||||
storage = create_storage(database_path)
|
||||
storage.set_user_model(42, "yandexgpt", "yandex")
|
||||
|
||||
self.assertEqual(storage.get_user_model(42, "yandex"), "yandexgpt")
|
||||
with closing(sqlite3.connect(database_path)) as connection:
|
||||
version = connection.execute(
|
||||
"PRAGMA user_version"
|
||||
).fetchone()[0]
|
||||
self.assertEqual(version, LATEST_SCHEMA_VERSION)
|
||||
|
||||
def test_default_models_are_injected(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
storage = create_storage(Path(directory) / "assistant.sqlite3")
|
||||
|
||||
self.assertEqual(
|
||||
storage.get_user_model(42, "local"),
|
||||
DEFAULT_MODELS["local"],
|
||||
)
|
||||
self.assertEqual(
|
||||
storage.get_user_model(42, "yandex"),
|
||||
DEFAULT_MODELS["yandex"],
|
||||
)
|
||||
|
||||
def test_provider_models_are_stored_separately(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
storage = AssistantStorage(Path(directory) / "assistant.sqlite3")
|
||||
storage = create_storage(Path(directory) / "assistant.sqlite3")
|
||||
|
||||
storage.set_user_model(42, "qwen3.5:9b", "local")
|
||||
storage.set_user_model(42, "yandexgpt", "yandex")
|
||||
@@ -244,7 +285,7 @@ class StorageTests(unittest.TestCase):
|
||||
|
||||
def test_memory_crud(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
storage = AssistantStorage(Path(directory) / "assistant.sqlite3")
|
||||
storage = create_storage(Path(directory) / "assistant.sqlite3")
|
||||
memory_id = storage.add_memory(42, "короткие ответы")
|
||||
|
||||
self.assertEqual(storage.list_memories(42)[0]["id"], memory_id)
|
||||
@@ -253,7 +294,7 @@ class StorageTests(unittest.TestCase):
|
||||
|
||||
def test_new_context_preserves_searchable_archive(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
storage = AssistantStorage(Path(directory) / "assistant.sqlite3")
|
||||
storage = create_storage(Path(directory) / "assistant.sqlite3")
|
||||
storage.add_conversation_exchange(
|
||||
42,
|
||||
100,
|
||||
|
||||
Reference in New Issue
Block a user