Add tests for table migration

Author: linuxdaemon

plugins/tell.py

@@ -78,29 +78,31 @@ def mark_read(self, now=None):
 
 @hook.on_start(priority=Priority.HIGHEST)
 def migrate_tables(db):
-    table = Table(
+    inspector = sa.inspect(db.bind)
+    if not inspector.has_table("tells"):
+        return
+
+    table = sa.Table(
         "tells",
         database.metadata,
-        Column("connection", String),
-        Column("sender", String),
-        Column("target", String),
-        Column("message", String),
-        Column("is_read", Boolean),
-        Column("time_sent", DateTime),
-        Column("time_read", DateTime),
+        autoload_with=db.bind,
     )
 
-    if not table.exists(db.bind):
-        return
-
-    if TellMessage.__table__.exists(db.bin):
+    if (
+        inspector.has_table(TellMessage.__tablename__)
+        and db.query(TellMessage).count() > 0
+    ):
         raise Exception(
-            f"Can't migrate table {table.name} to {TellMessage.__table__.name}, destination already exists"
+            f"Can't migrate table {table.name} to {TellMessage.__tablename__}, destination already exists"
         )
 
-    data = db.execute(table.select())
+    data = [dict(row) for row in db.execute(table.select())]
+    for item in data:
+        item["conn"] = item.pop("connection")
+
     db.bulk_insert_mappings(TellMessage, data, return_defaults=True)
     db.commit()
+
     table.drop(db.bind)
 
 

tests/plugin_tests/test_tell.py

@@ -1,8 +1,10 @@
 import datetime
 from unittest.mock import MagicMock, call, patch
 
+import sqlalchemy as sa
 from irclib.parser import Prefix
 
+from cloudbot.util import database
 from plugins import tell
 from tests.util.mock_conn import MockConn
 
@@ -18,6 +20,65 @@ def init_tables(mock_db):
     tell.load_ignores(session)
 
 
+def test_migrate_db(mock_db, freeze_time):
+    init_tables(mock_db)
+    session = mock_db.session()
+
+    tbl = sa.Table(
+        "tells",
+        database.metadata,
+        sa.Column("connection", sa.String),
+        sa.Column("sender", sa.String),
+        sa.Column("target", sa.String),
+        sa.Column("message", sa.String),
+        sa.Column("is_read", sa.Boolean),
+        sa.Column("time_sent", sa.DateTime),
+        sa.Column("time_read", sa.DateTime),
+    )
+
+    tbl.create(mock_db.engine)
+    mock_db.add_row(
+        tbl,
+        connection="conn",
+        sender="foo",
+        target="bar",
+        message="foobar",
+        is_read=False,
+        time_sent=datetime.datetime.now(),
+        time_read=None,
+    )
+
+    database.metadata.clear()
+
+    assert mock_db.get_data(tbl) == [
+        (
+            "conn",
+            "foo",
+            "bar",
+            "foobar",
+            False,
+            datetime.datetime(2019, 8, 22, 13, 14, 36),
+            None,
+        )
+    ]
+
+    tell.migrate_tables(session)
+
+    assert not sa.inspect(mock_db.engine).has_table(tbl.name)
+    assert mock_db.get_data(tell.TellMessage.__table__) == [
+        (
+            1,
+            "conn",
+            "foo",
+            "bar",
+            "foobar",
+            False,
+            datetime.datetime(2019, 8, 22, 13, 14, 36),
+            None,
+        )
+    ]
+
+
 def test_tellcmd(mock_db):
     init_tables(mock_db)
     session = mock_db.session()