diff --git a/test_tmdb.py b/test_tmdb.py
index 9740d04..e3b4ac3 100644
--- a/test_tmdb.py
+++ b/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):
Acting: Edward Norton, Brad Pitt, Helena Bonham Carter
taken from www.themoviedb.org
""")
+ 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()
diff --git a/tmdb/database.py b/tmdb/database.py
new file mode 100644
index 0000000..fdba0dd
--- /dev/null
+++ b/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 .
+'''
+
+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']
\ No newline at end of file
diff --git a/tmdb/tmdb.py b/tmdb/tmdb.py
index e798644..2aefa0f 100644
--- a/tmdb/tmdb.py
+++ b/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