From 97fbfe52ceb16bb39dca6e59e405ef66dda91860 Mon Sep 17 00:00:00 2001
From: Stephen Zhou <38493346+hyoban@users.noreply.github.com>
Date: Mon, 17 Aug 2026 17:42:26 +0800
Subject: [PATCH] test(knowledge-fs): cover icon background migration
 (WTA-2068)

---
 .../test_knowledge_fs_control_plane.py        |  5 +-
 .../test_knowledge_fs_icon_background.py      | 60 +++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 api/tests/unit_tests/migrations/test_knowledge_fs_icon_background.py

diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_control_plane.py b/api/tests/unit_tests/migrations/test_knowledge_fs_control_plane.py
index 73992379bad..3ddeb0f8beb 100644
--- a/api/tests/unit_tests/migrations/test_knowledge_fs_control_plane.py
+++ b/api/tests/unit_tests/migrations/test_knowledge_fs_control_plane.py
@@ -35,6 +35,7 @@ _TABLES = {
 }
 
 _FORBIDDEN_DEPENDENCIES = ("datasets", "documents", "dataset_permissions", "api_tokens")
+_LATER_COLUMNS_BY_TABLE = {"knowledge_fs_control_spaces": {"icon_background"}}
 
 _MODELS_BY_TABLE = {
     model.__tablename__: model
@@ -82,7 +83,9 @@ def test_upgrade_creates_only_independent_knowledge_fs_control_plane_tables() ->
     assert set(inspector.get_table_names()) >= _TABLES
     for table_name in _TABLES:
         migrated_columns = {column["name"] for column in inspector.get_columns(table_name)}
-        assert migrated_columns == set(_MODELS_BY_TABLE[table_name].__table__.columns.keys())
+        assert migrated_columns == (
+            set(_MODELS_BY_TABLE[table_name].__table__.columns.keys()) - _LATER_COLUMNS_BY_TABLE.get(table_name, set())
+        )
         dependency_text = " ".join(
             [
                 table_name,
diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_icon_background.py b/api/tests/unit_tests/migrations/test_knowledge_fs_icon_background.py
new file mode 100644
index 00000000000..7ee2ffa4d1e
--- /dev/null
+++ b/api/tests/unit_tests/migrations/test_knowledge_fs_icon_background.py
@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import sqlalchemy as sa
+from alembic.migration import MigrationContext
+from alembic.operations import Operations
+
+from models.knowledge_fs import KnowledgeFSControlSpace
+
+_MIGRATION_PATH = (
+    Path(__file__).resolve().parents[3]
+    / "migrations/versions/2026_08_17_1200-4f8b2c7d9e10_add_knowledge_fs_icon_background.py"
+)
+
+
+def _load_migration():
+    spec = importlib.util.spec_from_file_location(_MIGRATION_PATH.stem, _MIGRATION_PATH)
+    if spec is None or spec.loader is None:
+        raise RuntimeError("failed to load KnowledgeFS icon-background migration")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _run(module: object, engine: sa.Engine, step: str) -> None:
+    with engine.begin() as connection:
+        operations = Operations(MigrationContext.configure(connection))
+        original_op = module.op
+        module.op = operations
+        try:
+            getattr(module, step)()
+        finally:
+            module.op = original_op
+
+
+def test_icon_background_migration_matches_model_and_merges_heads() -> None:
+    engine = sa.create_engine("sqlite:///:memory:")
+    with engine.begin() as connection:
+        connection.exec_driver_sql("CREATE TABLE knowledge_fs_control_spaces (id VARCHAR(36) PRIMARY KEY)")
+    migration = _load_migration()
+
+    _run(migration, engine, "upgrade")
+
+    columns = {
+        column["name"]: column for column in sa.inspect(engine).get_columns(KnowledgeFSControlSpace.__tablename__)
+    }
+    icon_background = columns["icon_background"]
+    assert icon_background["nullable"] is False
+    assert icon_background["type"].length == 7
+    assert "F0F9FF" in icon_background["default"]
+    assert migration.revision == "4f8b2c7d9e10"
+    assert migration.down_revision == ("9d4e6f8a1b2c", "56124e050600")
+
+    _run(migration, engine, "downgrade")
+
+    assert "icon_background" not in {
+        column["name"] for column in sa.inspect(engine).get_columns(KnowledgeFSControlSpace.__tablename__)
+    }