Moved database to own file

test_tmdb.py

@@ -2,6 +2,9 @@
 import unittest
 from tmdb.tmdb_api import Movie
 from tmdb.tmdb import TmdbBot
+from tmdb.database import Database
+from sqlalchemy import create_engine
+
 
 class TestTmdbMethods(unittest.TestCase):
     ### TMDB API
@@ -43,5 +46,13 @@ class TestTmdbMethods(unittest.TestCase):
                            <p>Acting: Edward Norton, Brad Pitt, Helena Bonham Carter</p>
                            <p>taken from www.themoviedb.org</p>""")
 
+    def test_database_language(self):
+        engine = create_engine('sqlite:///test.db', echo = True)
+        db = Database(engine)
+        db.set_language('@testuser:example.com', 'de')
+        self.assertEqual(str(db.get_language('@testuser:example.com')), 'de')
+        db.set_language('@testuser:example.com', 'en')
+        self.assertEqual(str(db.get_language('@testuser:example.com')), 'en')
+
 if __name__ == '__main__':
     unittest.main()

tmdb/database.py

@@ -0,0 +1,44 @@
+'''
+This file is part of tmdb-bot.
+
+tmdb-bot is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General  Public License version 3 as published by
+the Free Software Foundation.
+
+tmdb-bot is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General  Public License for more details.
+
+You should have received a copy of the GNU Affero General  Public License
+along with tmdb-bot.  If not, see <https://www.gnu.org/licenses/>.
+'''
+
+from sqlalchemy import (Column, String, Integer, ForeignKey, Table, MetaData,
+                        select, and_)
+from sqlalchemy.engine.base import Engine
+
+class Database:
+    db: Engine
+
+    def __init__(self, db: Engine) -> None:
+        self.db = db
+        
+        meta = MetaData()
+        meta.bind = db
+
+        self.language = Table("tmdb_language", meta,
+                              Column("id", Integer, primary_key=True, autoincrement=True),
+                              Column("user_id", String(255), nullable=False),
+                              Column("language", String(255), nullable=False),)
+        meta.create_all(db)
+        
+    def set_language(self, user_id, language):
+        with self.db.begin() as tx:
+            tx.execute(self.language.delete().where(self.language.c.user_id == user_id))
+            tx.execute(self.language.insert().values(user_id=user_id, language=language))
+            
+    def get_language(self, user_id):
+        rows = self.db.execute(select([self.language.c.language])
+                                   .where(self.language.c.user_id == user_id))
+        return rows.fetchone()['language']

tmdb/tmdb.py

@@ -21,35 +21,8 @@ from maubot import Plugin, MessageEvent
 from maubot.handlers import command
 
 from tmdb.tmdb_api import Movie
+from tmdb.database import Database
 
-from sqlalchemy import (Column, String, Integer, ForeignKey, Table, MetaData,
-                        select, and_)
-from sqlalchemy.engine.base import Engine
-
-class Database:
-    db: Engine
-
-    def __init__(self, db: Engine) -> None:
-        self.db = db
-        
-        meta = MetaData()
-        meta.bind = db
-
-        self.language = Table("tmdb_language", meta,
-                              Column("id", Integer, primary_key=True, autoincrement=True),
-                              Column("user_id", String(255), nullable=False),
-                              Column("language", String(255), nullable=False),)
-        meta.create_all(db)
-        
-    def set_language(self, user_id, language):
-        with self.db.begin() as tx:
-            tx.execute(self.language.delete().where(self.language.c.user_id == user_id))
-            tx.execute(self.language.insert().values(user_id=user_id, language=language))
-            
-    def get_language(self, user_id):
-        rows = self.db.execute(select([self.language.c.language])
-                                   .where(self.language.c.user_id == user_id))
-        return rows.fetchone()
 
 class TmdbBot(Plugin):
     db: Database