From 75d2682458287e039efce1c1f9de5b43252f299c Mon Sep 17 00:00:00 2001 From: lomion Date: Tue, 23 Jun 2020 17:34:34 +0200 Subject: [PATCH] Moved database to own file --- test_tmdb.py | 11 +++++++++++ tmdb/database.py | 44 ++++++++++++++++++++++++++++++++++++++++++++ tmdb/tmdb.py | 29 +---------------------------- 3 files changed, 56 insertions(+), 28 deletions(-) create mode 100644 tmdb/database.py 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