"""Data groups: how one is resolved, the startup sweep, deleting one, and upgrading. What a group *isolates* is `test_data_group_isolation.py`; how a chat is pinned to one is `test_data_group_chat_pin.py`. This file is the machinery underneath both: the resolution order, which lives in one function and must keep living there, and the upgrade from a database that has never heard of groups. """ from __future__ import annotations import pytest from sqlalchemy import inspect, select, text from lembas.db.migrations import sync_schema from lembas.db.models import ( DEFAULT_GROUP, Chat, Connection, DataGroup, Memory, Model, Note, Persona, User, ) from lembas.db.session import get_engine from lembas.services import data_groups, personas, settings_store from lembas.services.crypto import encrypt @pytest.fixture def owner(db, registered) -> User: return db.scalars(select(User).order_by(User.created_at)).first() @pytest.fixture def reader(db, registered) -> User: """A second account, not an administrator, so permissions actually apply.""" user = User(email="sam@shire.test", name="Sam", password_hash="x", role="user") db.add(user) db.commit() return user @pytest.fixture def two(db) -> tuple[Connection, Connection]: """A local connection in the default group and a hosted one in its own.""" db.add(DataGroup(id="hosted", name="Hosted")) local = Connection(name="Local", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt("")) cloud = Connection( name="Cloud", base_url="http://127.0.0.1:2", api_key_encrypted=encrypt(""), data_group_id="hosted", ) db.add_all([local, cloud]) db.flush() db.add_all( [ Model(connection_id=local.id, model_id="local-model", position=0), Model(connection_id=cloud.id, model_id="cloud-model", position=1), ] ) db.commit() return local, cloud def _grant_manage(db, user: User) -> None: settings_store.update(db, {"default_permissions": {data_groups.PERMISSION: True}}) # --- Resolution -------------------------------------------------------------------- def test_a_connection_with_no_group_is_in_the_default_one(db, two, reader): local, _ = two assert data_groups.for_connection(db, reader, local.id) == DEFAULT_GROUP def test_the_administrators_choice_applies_to_everybody(db, two, reader): _, cloud = two assert data_groups.for_connection(db, reader, cloud.id) == "hosted" assert data_groups.for_connection(db, None, cloud.id) == "hosted" def test_a_personal_mapping_needs_the_permission(db, two, reader): """Stored and ignored without `data.manage`: taking the permission away puts a person back on the instance's arrangement without anybody clearing anything.""" local, _ = two reader.settings_json = {data_groups.SETTING_KEY: {local.id: "hosted"}} db.commit() assert data_groups.for_connection(db, reader, local.id) == DEFAULT_GROUP _grant_manage(db, reader) assert data_groups.for_connection(db, reader, local.id) == "hosted" def test_a_mapping_to_somebody_elses_personal_group_is_ignored(db, two, reader, owner): local, _ = two db.add(DataGroup(id="theirs", name="Theirs", owner_id=owner.id)) reader.settings_json = {data_groups.SETTING_KEY: {local.id: "theirs"}} db.commit() _grant_manage(db, reader) assert data_groups.for_connection(db, reader, local.id) == DEFAULT_GROUP def test_a_connection_naming_a_deleted_group_falls_back_to_the_default(db, two, reader): _, cloud = two cloud.data_group_id = "gone" db.commit() assert data_groups.for_connection(db, reader, cloud.id) == DEFAULT_GROUP def test_a_pair_without_a_connection_resolves_through_the_model(db, two, reader): assert data_groups.for_pair(db, reader, "cloud-model") == "hosted" assert data_groups.for_pair(db, reader, "local-model") == DEFAULT_GROUP def test_a_chat_keeps_the_group_it_was_stamped_with(db, two, reader): """Derived once, then read. A model moved afterwards does not carry the chat.""" _, cloud = two chat = Chat(user_id=reader.id, model_id="cloud-model", connection_id=cloud.id) db.add(chat) db.commit() assert data_groups.for_chat(db, chat) == "hosted" db.commit() cloud.data_group_id = DEFAULT_GROUP db.commit() assert data_groups.for_chat(db, chat) == "hosted" def test_usable_groups_are_the_instances_and_ones_own(db, two, reader, owner): db.add_all( [ DataGroup(id="mine", name="Mine", owner_id=reader.id), DataGroup(id="not-mine", name="Not mine", owner_id=owner.id), ] ) db.commit() ids = {group.id for group in data_groups.usable(db, reader)} assert ids == {DEFAULT_GROUP, "hosted", "mine"} def test_one_group_means_there_is_nothing_to_choose(db, reader): assert data_groups.several(db, reader) is False db.add(DataGroup(id="second", name="Second")) db.commit() assert data_groups.several(db, reader) is True # --- Deleting ------------------------------------------------------------------------ def test_the_default_group_cannot_be_deleted(db): with pytest.raises(ValueError, match="default group"): data_groups.delete(db, data_groups.ensure_default(db)) def test_a_group_with_records_in_it_cannot_be_deleted(db, reader): group = DataGroup(id="busy", name="Busy") db.add(group) db.add(Note(owner_id=reader.id, title="n", body="b", data_group_id="busy")) db.commit() with pytest.raises(ValueError, match="1 notes"): data_groups.delete(db, group) def test_a_group_a_connection_is_in_cannot_be_deleted(db, two): with pytest.raises(ValueError, match="connections"): data_groups.delete(db, data_groups.get(db, "hosted")) def test_deleting_an_empty_group_clears_every_mapping_to_it(db, reader): group = DataGroup(id="empty", name="Empty") db.add(group) reader.settings_json = {data_groups.SETTING_KEY: {"some-connection": "empty"}} db.commit() data_groups.delete(db, group) db.refresh(reader) assert data_groups.personal_map(reader) == {} assert data_groups.get(db, "empty") is None # --- The sweep ------------------------------------------------------------------------- def test_the_sweep_puts_old_rows_in_the_default_group(db, reader): db.add(Memory(owner_id=reader.id, content="old")) db.commit() assert db.scalar(select(Memory)).data_group_id is None data_groups.sweep_unassigned(db) assert db.scalar(select(Memory)).data_group_id == DEFAULT_GROUP def test_the_sweep_gives_a_chat_its_models_group(db, two, reader): """Not simply the default: a chat some path created without stamping one belongs where its model is, or its next turn is refused as a moved chat.""" chat = Chat(user_id=reader.id, model_id="cloud-model") db.add(chat) db.commit() data_groups.sweep_unassigned(db) db.refresh(chat) assert chat.data_group_id == "hosted" def test_a_row_the_sweep_has_not_reached_still_counts_as_default(db, reader): db.add(Note(owner_id=reader.id, title="old", body="")) db.commit() found = db.scalars(select(Note).where(data_groups.condition(Note, DEFAULT_GROUP))).all() assert [note.title for note in found] == ["old"] # --- Personalities: namespaced, because the constraint cannot change ------------------ def test_the_default_group_keeps_the_bare_model_id(): assert personas.key_for("gpt-oss", DEFAULT_GROUP) == "gpt-oss" assert personas.key_for("gpt-oss", None) == "gpt-oss" def test_another_group_gets_its_own_key_and_splits_back(): key = personas.key_for("gpt-oss", "hosted") assert key != "gpt-oss" assert personas.split_key(key) == ("gpt-oss", "hosted") assert personas.split_key("gpt-oss") == ("gpt-oss", DEFAULT_GROUP) def test_a_group_falls_back_to_the_administrators_default(db, reader): """The admin default is keyed bare and reaches every group, until the model has written one of its own with this person in that group.""" db.add(Persona(model_key="gpt-oss", owner_id=None, content="the default")) db.commit() key = personas.key_for("gpt-oss", "hosted") assert personas.block(db, key, reader) == "the default" personas.write(db, model_key=key, owner=reader, content="hosted self", author="model") assert personas.block(db, key, reader) == "hosted self" assert personas.block(db, "gpt-oss", reader) == "the default" # --- Upgrading from 1.9.1 ----------------------------------------------------------------- # What 1.10.0 added: one table, and one column on each of these. Taken from the # models rather than invented, and `test_the_recorded_shape_is_still_real` # below is what stops the list rotting. NEW_TABLE = "data_groups" GROUPED_TABLES = ( "chats", "connections", "knowledge_bases", "memories", "notes", "reports", "schedules", "skills", ) def _rollback_to_1_9_1(engine) -> None: with engine.begin() as connection: connection.execute(text(f"DROP TABLE IF EXISTS {NEW_TABLE}")) for table in GROUPED_TABLES: connection.execute(text(f"ALTER TABLE {table} DROP COLUMN data_group_id")) def test_the_recorded_shape_is_still_real(): from lembas.db.base import Base grouped = { table.name for table in Base.metadata.sorted_tables if "data_group_id" in table.columns } assert grouped == set(GROUPED_TABLES) def test_a_1_9_1_database_with_data_upgrades(db, reader): engine = get_engine() db.add(Memory(owner_id=reader.id, content="kept")) db.add(Persona(model_key="gpt-oss", owner_id=reader.id, content="mine")) db.commit() db.close() _rollback_to_1_9_1(engine) assert "data_group_id" not in {c["name"] for c in inspect(engine).get_columns("memories")} changes = sync_schema(engine) assert any(NEW_TABLE in change for change in changes) for table in GROUPED_TABLES: assert "data_group_id" in {c["name"] for c in inspect(engine).get_columns(table)} # A nullable column arrives empty; the sweep is what files it. memory = db.scalar(select(Memory)) assert memory.data_group_id is None data_groups.sweep_unassigned(db) db.refresh(memory) assert memory.data_group_id == DEFAULT_GROUP # The person's personality was keyed bare, and the default group still reads it. user = db.get(User, reader.id) assert personas.block(db, personas.key_for("gpt-oss", DEFAULT_GROUP), user) == "mine" assert data_groups.get(db, DEFAULT_GROUP) is not None