bot: test peewee sqlite3

1 files changed, 24 insertions, 27 deletions

diff --git a/bot.py b/bot.py
index 79fba5d..a059c9c 100644
--- a/bot.py
+++ b/bot.py
@@ -61,6 +61,7 @@ from config import (
 )
 from mastodon import Mastodon
 
+from dbstore.peewee_store import db, User
 
 # Enable logging
 logging.basicConfig(
@@ -72,39 +73,35 @@ logger = logging.getLogger(__name__)
 @dataclass
 class FediLoginCallbackUpdate:
     code: str
-    state: int
+    state: str
 
 
 class FediLoginCallbackContext(CallbackContext[ExtBot, dict, dict, dict]):
-    @classmethod
-    def from_update(
-            cls,
-            update: object,
-            application: "Application",
-    ) -> "FediLoginCallbackContext":
-        if isinstance(update, FediLoginCallbackUpdate):
-            return cls(application=application, user_id=update.state)
-        return super().from_update(update, application)
+    pass
 
 
 async def process_oauth_login_callback(update: FediLoginCallbackUpdate, context: FediLoginCallbackContext) -> None:
-    # query client_id and client_secret from in memory database
-    client_id = ""
-    client_secret = ""
-    home_instance = ""
-
-    mastodon_client = Mastodon(client_id=client_id, client_secret=client_secret, api_base_url=home_instance)
-    access_token = mastodon_client.log_in(
-        code=update.code,
-        redirect_uri="{}{}".format(BOT_DOMAIN, FEDI_LOGIN_CALLBACK_URL),
-        scopes=BOT_SCOPE
-    )
+    state = update.state
+
+    with db.connection_context():
+        user = User.get(User.state == state)
 
-    # TODO
-    # save access_token to database
+        client_id = user.client_id
+        client_secret = user.client_secret
+        home_instance = user.home_instance
+
+        if len(user.access_key) == 0:
+            mastodon_client = Mastodon(client_id=client_id, client_secret=client_secret, api_base_url=home_instance)
+            access_token = mastodon_client.log_in(
+                code=update.code,
+                redirect_uri="{}{}".format(BOT_DOMAIN, FEDI_LOGIN_CALLBACK_URL),
+                scopes=BOT_SCOPE
+            )
+            user.access_key = access_token
+            user.save()
 
-    text = "You have successfully logged in to your Mastodon account!"
-    await context.bot.send_message(chat_id=update.state, text=text)
+        text = "You have successfully logged in to your Mastodon account!"
+        await context.bot.send_message(chat_id=user.telegram_user_id, text=text)
 
 
 async def main() -> None:
@@ -112,7 +109,7 @@ async def main() -> None:
     # Here we set updater to None because we want our custom webhook server to handle the updates
     # and hence we don't need an Updater instance
     application = (
-        Application.builder().token(BOT_TOKEN).context_types(context_types).build()
+        Application.builder().updater(None).token(BOT_TOKEN).context_types(context_types).build()
     )
 
     checkin_handler = ConversationHandler(
@@ -171,7 +168,7 @@ async def main() -> None:
         """
         try:
             code = request.query_params["code"]
-            state = int(request.query_params.get("state"))
+            state = request.query_params.get("state")
         except KeyError:
             return PlainTextResponse(
                 status_code=HTTPStatus.BAD_REQUEST,