diff --git a/bot.py b/bot.py index 7dbdae8..49028e7 100644 --- a/bot.py +++ b/bot.py @@ -1,26 +1,25 @@ import asyncio import base64 -import boto3 import os import spotipy import logging import re import json import telegram -from enum import Enum from telegram import Update, LinkPreviewOptions from telegram.ext import ( Application, CommandHandler, MessageHandler, filters, - CallbackContext, Defaults, + ContextTypes, ) from spotipy.oauth2 import SpotifyOAuth, CacheHandler from spotipy.exceptions import SpotifyException import urllib.parse - +from custom_callback_contex import CustomCallbackContext +from storage.storage_interface import BotState if logging.getLogger().hasHandlers(): logging.getLogger().setLevel(logging.INFO) @@ -29,7 +28,6 @@ logger = logging.getLogger() # Constants and configurations - defaults = Defaults( link_preview_options=LinkPreviewOptions(show_above_text=False, is_disabled=True), disable_notification=True, @@ -38,20 +36,6 @@ SPOTIFY_CLIENT_SECRET = os.getenv("SPOTIFY_CLIENT_SECRET") SPOTIFY_REDIRECT_URI = os.getenv("SPOTIFY_REDIRECT_URI") spotify_link_pattern = r"https://open\.spotify\.com/track/([a-zA-Z0-9]+)" -# Dynamodb -dynamodb = boto3.resource("dynamodb", region_name="us-east-1") - -bot_table = dynamodb.Table(os.getenv("BOT_TABLE")) -credentials_table = dynamodb.Table(os.getenv("CREDENTIALS_TABLE")) - - -# Enums for bot states -class BotState(Enum): - AWAITING_PLAYLIST_IMAGE = "awaiting_playlist_image" - CHANGING_PLAYLIST_NAME = "changing_playlist_name" - CREATING_PLAYLIST = "creating_playlist" - CHANGING_PLAYLIST_IMAGE = "changing_playlist_image" - NO_STATE = None def load_html_file(file_name): @@ -59,13 +43,13 @@ def load_html_file(file_name): return file.read() -def handle_spotify_auth(state, code): +def handle_spotify_auth(state, code, storage): state_decoded = urllib.parse.unquote(state) state_info = json.loads(state_decoded) chat_id = state_info.get("chat_id") user_id = state_info.get("user_id") - sp_oauth = get_sp_oauth(chat_id, user_id) + sp_oauth = get_sp_oauth(chat_id, user_id, storage) token_info = sp_oauth.get_access_token(code) if token_info: @@ -81,158 +65,42 @@ def handle_spotify_auth(state, code): # return {"statusCode": 200, "body": json.dumps(event, indent=2)} -# ----------------------------------------------- -# DynamoDB Utility Functions -# ----------------------------------------------- -def save_current_state(chat_id, state_key: BotState): - # Saves the current bot state for a given chat_id in DynamoDB. - try: - user_id = get_user_id_from_chat_id(chat_id) - if user_id is None: - user_id = get_user_id_from_channel_credentials(chat_id) - - if state_key is BotState.NO_STATE: - response = bot_table.update_item( - Key={"chat_id": str(chat_id)}, - UpdateExpression="REMOVE current_state", - ReturnValues="UPDATED_NEW", - ) - else: - response = bot_table.update_item( - Key={"chat_id": str(chat_id)}, - UpdateExpression="SET current_state = :state, user_id = :uid", - ExpressionAttributeValues={ - ":state": state_key.value, - ":uid": user_id, - }, - ReturnValues="UPDATED_NEW", - ) - logging.info(f"Updated state in DynamoDB: {response}") - return response - except Exception as e: - logging.error(f"Error saving current state to DynamoDB: {e}") - - -def get_current_state(chat_id): - # Retrieves the current bot state for a given chat_id from DynamoDB. - try: - response = bot_table.get_item(Key={"chat_id": str(chat_id)}) - if "Item" in response and "current_state" in response["Item"]: - state_value = response["Item"]["current_state"] - return BotState(state_value) - return None - except Exception as e: - logging.error(f"Error retrieving current state from DynamoDB: {e}") - return None - - -def save_playlist_to_dynamodb(chat_id, playlist_id): - try: - bot_table.put_item( - Item={ - "chat_id": str(chat_id), - "playlist_id": playlist_id, - "user_id": get_user_id_from_chat_id(chat_id), - } - ) - except Exception as e: - logging.error(f"Error saving to DynamoDB: {e}") - - -def get_playlist_from_dynamodb(chat_id): - try: - response = bot_table.get_item(Key={"chat_id": str(chat_id)}) - if "Item" in response: - playlist_id = response["Item"].get("playlist_id") - return playlist_id - return None - except Exception as e: - logging.error(f"Error retrieving from DynamoDB: {e}") - return None - - -def get_user_id_from_chat_id(chat_id): - try: - response = bot_table.get_item(Key={"chat_id": str(chat_id)}) - if "Item" in response and "user_id" in response["Item"]: - return response["Item"]["user_id"] - else: - logging.info( - f"get_user_id_from_chat_id: No user_id found for chat_id: {chat_id}: {response}" - ) - return None - except Exception as e: - logging.error( - f"Error retrieving user ID from DynamoDB for chat_id: {chat_id}, error: {e}" - ) - return None - - -def get_user_id_from_channel_credentials(chat_id): - try: - response = credentials_table.get_item(Key={"chat_id": str(chat_id)}) - if "Item" in response and "user_id" in response["Item"]: - return response["Item"]["user_id"] - else: - print( - f"get_user_id_from_channel_credentials: No user_id found for chat_id: {chat_id}: {response}" - ) - return None - except Exception as e: - print( - f"Error retrieving user ID from DynamoDB for chat_id: {chat_id}, error: {e}" - ) - return None - - class DynamoCredentialsCache(CacheHandler): """ A cache handler that stores OAuth credentials in a Dynamo bot_table called 'ChannelCredentials' that has a primary key of chat_id. """ - def __init__(self, chat_id, user_id): + def __init__(self, chat_id, user_id, storage): self.chat_id = chat_id self.user_id = user_id + self.storage = storage def get_cached_token(self): try: - response = credentials_table.get_item(Key={"chat_id": str(self.chat_id)}) - if "Item" in response: - return response["Item"] - return None + return self.storage.get_cached_token(self.chat_id) except Exception as e: - logging.error(f"Error retrieving from DynamoDB: {e}") + logging.error(f"Error retrieving cached token from DynamoDB: {e}") raise def save_token_to_cache(self, token_info): try: - credentials_table.put_item( - Item={ - "chat_id": str(self.chat_id), - "user_id": self.user_id, - **token_info, - } - ) - bot_table.update_item( - Key={"chat_id": str(self.chat_id)}, - UpdateExpression="SET user_id = :uid", - ExpressionAttributeValues={":uid": str(self.user_id)}, - ) + self.storage.save_token_to_cache(self.chat_id, self.user_id, token_info) except Exception as e: - logging.error(f"Error saving to DynamoDB: {e}") + logging.error(f"Error saving token to DynamoDB: {e}") raise # ----------------------------------------------- # Spotify Utility Functions # ----------------------------------------------- -def get_sp_oauth(chat_id, user_id): +def get_sp_oauth(chat_id, user_id, storage): + # TODO: Pass Storage Param return SpotifyOAuth( SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, SPOTIFY_REDIRECT_URI, - cache_handler=DynamoCredentialsCache(chat_id, user_id), + cache_handler=DynamoCredentialsCache(chat_id, user_id, storage), scope="playlist-modify-public ugc-image-upload", ) @@ -272,23 +140,23 @@ def create_spotify_playlist(playlist_name, sp_oauth): # ----------------------------------------------- # Telegram Message Handlers # ----------------------------------------------- -async def handle_playlist_image(update: Update, context: CallbackContext) -> None: +async def handle_playlist_image(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id user_id = update.effective_user.id - user_id_table = get_user_id_from_chat_id(chat_id) + user_id_table = context.storage().get_user_id_from_chat_id(chat_id) if user_id_table != str(user_id): await update.message.reply_text( "You are not authorized for this playlist process." ) return - current_state = get_current_state(chat_id) + current_state = context.storage().get_current_state(chat_id) if ( current_state == BotState.AWAITING_PLAYLIST_IMAGE or current_state == BotState.CHANGING_PLAYLIST_IMAGE ): photo = update.message.photo[-1] await update.message.reply_text("Processing your image, please wait...") - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if not playlist_id: await update.message.reply_text("No playlist found for this chat.") return @@ -304,7 +172,7 @@ async def handle_playlist_image(update: Update, context: CallbackContext) -> Non return async def upload_image(): - sp_oauth = get_sp_oauth(chat_id, user_id) + sp_oauth = get_sp_oauth(chat_id, user_id, context.storage()) sp = spotipy.Spotify(auth_manager=sp_oauth) sp.playlist_upload_cover_image(playlist_id, base64_image) @@ -320,7 +188,7 @@ async def upload_image(): ), ) - save_current_state(chat_id, BotState.NO_STATE) + context.storage().save_current_state(chat_id, BotState.NO_STATE) except asyncio.TimeoutError: logging.exception("Error Timeout: TimeoutError") await update.message.reply_text("Image upload timed out. Please try again.") @@ -329,14 +197,19 @@ async def upload_image(): await update.message.reply_text(f"An error occurred: {e}") -async def handle_playlist_name(update: Update, context: CallbackContext) -> None: +async def handle_playlist_name(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id user_id = update.effective_user.id - user_id_bot_table = get_user_id_from_chat_id(chat_id) - user_id_credentials_table = get_user_id_from_channel_credentials(chat_id) + user_id_bot_table = context.storage().get_user_id_from_chat_id(chat_id) + user_id_credentials_table = context.storage().get_user_id_from_channel_credentials( + chat_id + ) + logger.info(f"user id table: {user_id_credentials_table}, user id: {user_id} ") - current_state = get_current_state(chat_id) - sp_oauth = get_sp_oauth(chat_id, user_id) + current_state = context.storage().get_current_state(chat_id) + logging.info(f"the state here-pre is: {current_state}, {type(current_state)}") + sp_oauth = get_sp_oauth(chat_id, user_id, context.storage()) + logging.info(f"here is sp oauth: {sp_oauth}") if current_state == BotState.CHANGING_PLAYLIST_NAME: if user_id_bot_table != str(user_id): @@ -344,7 +217,7 @@ async def handle_playlist_name(update: Update, context: CallbackContext) -> None "Please click on authorize link before entering playlist name." ) return - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if playlist_id: new_name = update.message.text.strip() if change_spotify_playlist_name(playlist_id, new_name, sp_oauth): @@ -353,12 +226,13 @@ async def handle_playlist_name(update: Update, context: CallbackContext) -> None await update.message.reply_text( "Failed to change the playlist name. Please try again later." ) - save_current_state(chat_id, BotState.NO_STATE) + context.storage().save_current_state(chat_id, BotState.NO_STATE) else: await update.message.reply_text( "Make sure you have a playlist created before changing the name." ) - elif current_state == BotState.CREATING_PLAYLIST: + elif current_state is BotState.CREATING_PLAYLIST: + logging.info(f"the state here is: {current_state}") if str(user_id) != user_id_credentials_table: await update.message.reply_text( "You are not authorized for this playlist process." @@ -368,8 +242,10 @@ async def handle_playlist_name(update: Update, context: CallbackContext) -> None playlist_name = update.message.text.strip() try: playlist_id = create_spotify_playlist(playlist_name, sp_oauth) - save_playlist_to_dynamodb(chat_id, playlist_id) - save_current_state(chat_id, BotState.AWAITING_PLAYLIST_IMAGE) + context.storage().save_playlist_to_dynamodb(chat_id, playlist_id) + context.storage().save_current_state( + chat_id, BotState.AWAITING_PLAYLIST_IMAGE + ) await update.message.reply_text( f"Created new playlist: {playlist_name}. " "Now please send me a cool image to set as your playlist cover 😎." @@ -384,22 +260,22 @@ async def handle_playlist_name(update: Update, context: CallbackContext) -> None return -async def handle_spotify_links(update: Update, context: CallbackContext) -> None: +async def handle_spotify_links(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id user_id = update.effective_user.id message_text = update.message.text match = re.search(spotify_link_pattern, message_text) - current_state = get_current_state(chat_id) + current_state = context.storage().get_current_state(chat_id) if current_state == BotState.CREATING_PLAYLIST: await update.message.reply_text( "You are in the process of creating a playlist. " "Please wait until it's done before sending links." ) return - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if match and playlist_id: track_id = match.group(1) - sp_oauth = get_sp_oauth(chat_id, user_id) + sp_oauth = get_sp_oauth(chat_id, user_id, context.storage()) if add_track_to_spotify_playlist(playlist_id, track_id, sp_oauth): await update.message.set_reaction("👍") else: @@ -416,31 +292,31 @@ async def handle_spotify_links(update: Update, context: CallbackContext) -> None # ----------------------------------------------- # Telegram Command Handlers # ----------------------------------------------- -async def change_playlist_image(update: Update, context: CallbackContext) -> None: +async def change_playlist_image(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if playlist_id: await update.message.reply_text("Please send the new image for your playlist:") - save_current_state(chat_id, BotState.CHANGING_PLAYLIST_IMAGE) + context.storage().save_current_state(chat_id, BotState.CHANGING_PLAYLIST_IMAGE) else: await update.message.reply_text( "No playlist found for this chat. Create one with /createplaylist." ) -async def change_playlist_name(update: Update, context: CallbackContext) -> None: +async def change_playlist_name(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if playlist_id: await update.message.reply_text("Please enter the new name for your playlist:") - save_current_state(chat_id, BotState.CHANGING_PLAYLIST_NAME) + context.storage().save_current_state(chat_id, BotState.CHANGING_PLAYLIST_NAME) else: await update.message.reply_text( "No playlist found for this chat. Create one with /createplaylist." ) -async def create_playlist(update: Update, context: CallbackContext) -> bool: +async def create_playlist(update: Update, context: CustomCallbackContext) -> bool: chat_id = update.effective_chat.id user_id = update.effective_user.id @@ -448,9 +324,9 @@ async def create_playlist(update: Update, context: CallbackContext) -> bool: state_encoded = json.dumps(state_info) state_url_safe = urllib.parse.quote(state_encoded) - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) - current_state = get_current_state(chat_id) + current_state = context.storage().get_current_state(chat_id) if playlist_id: await update.message.reply_text( @@ -463,8 +339,8 @@ async def create_playlist(update: Update, context: CallbackContext) -> bool: "You are already in the process of creating a playlist." ) return False - save_current_state(chat_id, BotState.CREATING_PLAYLIST) - sp_oauth = get_sp_oauth(chat_id, user_id) + context.storage().save_current_state(chat_id, BotState.CREATING_PLAYLIST) + sp_oauth = get_sp_oauth(chat_id, user_id, context.storage()) token_info = sp_oauth.cache_handler.get_cached_token() if sp_oauth.validate_token(token_info) is None: auth_url = sp_oauth.get_authorize_url(state=state_url_safe) @@ -476,7 +352,7 @@ async def create_playlist(update: Update, context: CallbackContext) -> bool: return True -async def help_command(update: Update, context: CallbackContext) -> None: +async def help_command(update: Update, context: CustomCallbackContext) -> None: help_text = ( "Here are the commands you can use:\n" "/start - Start interacting with the bot\n" @@ -492,13 +368,12 @@ async def help_command(update: Update, context: CallbackContext) -> None: await update.message.reply_text(help_text) -async def reset_playlist(update: Update, context: CallbackContext) -> None: +async def reset_playlist(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id # Delete the playlist entry from DynamoDB try: - bot_table.delete_item(Key={"chat_id": str(chat_id)}) - except Exception as e: - logging.error(f"Error deleting from DynamoDB: {e}") + context.storage().delete_item(context.storage().bot_table, chat_id) + except Exception: await update.message.reply_text("Failed to reset the playlist in the database.") return await update.message.reply_text( @@ -507,9 +382,9 @@ async def reset_playlist(update: Update, context: CallbackContext) -> None: ) -async def send_playlist_link(update: Update, context: CallbackContext) -> None: +async def send_playlist_link(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if playlist_id: playlist_url = f"https://open.spotify.com/playlist/{playlist_id}" await update.message.reply_text( @@ -523,18 +398,46 @@ async def send_playlist_link(update: Update, context: CallbackContext) -> None: await update.message.reply_text("No playlist found for this chat.") -async def start(update: Update, context: CallbackContext) -> None: +async def start(update: Update, context: CustomCallbackContext) -> None: await update.message.reply_text( "Hiya! I'm your Spotify Skunk bot 🦨. /createplaylist " "to add songs to your playlist!" ) -async def unlink_credentials(update: Update, context: CallbackContext) -> None: +async def unlink_credentials(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id try: - credentials_table.delete_item(Key={"chat_id": str(chat_id)}) - bot_table.delete_item(Key={"chat_id": str(chat_id)}) + # Check if the item exists in the credentials table + if not context.storage().check_item_exists("credentials_table", chat_id): + logging.error(f"Credentials not found for chat_id {chat_id}") + await update.message.reply_text( + "Failed to unlink your Spotify credentials: credentials not found." + ) + return + + # Check if the item exists in the bot table + + if not context.storage().check_item_exists("bot_table", chat_id): + logging.error(f"Bot data not found for chat_id {chat_id}") + await update.message.reply_text( + "Failed to unlink your Spotify credentials: bot data not found." + ) + return + + # Log the deletion attempt + logging.info( + f"Attempting to delete item from credentials table for chat_id {chat_id}" + ) + credentials_delete_response = context.storage().delete_item( + "credentials_table", chat_id + ) + logging.info(f"Credentials delete_item response: {credentials_delete_response}") + + logging.info(f"Attempting to delete item from bot table for chat_id {chat_id}") + bot_delete_response = context.storage().delete_item("bot_table", chat_id) + logging.info(f"Bot delete_item response: {bot_delete_response}") + await update.message.reply_text( "Your Spotify credentials have been unlinked successfully." ) @@ -545,9 +448,34 @@ async def unlink_credentials(update: Update, context: CallbackContext) -> None: ) -def build_application(token): +# TODO: Fix so it Only Works for You and in Group Chats +async def debug_stuff(update: Update, context: CustomCallbackContext): + chat_id = update.effective_chat.id + debug = "\n".join( + [ + f"Owning User ID: `{context.storage().get_user_id_from_chat_id(chat_id)}`", + f"Chat ID: `{context.storage().get_current_state(chat_id)}`", + f"Playlist: `{context.storage().get_playlist_from_dynamodb(chat_id)}`", + f"User ID from Channel Credentials: `{context.storage().get_user_id_from_channel_credentials(chat_id)}`", + ] + ) + await update.message.reply_text(debug) + + +def build_application(token, storage): logger.info(f"token: {token}") - application = Application.builder().token(token).defaults(defaults).build() + + # Set the custom context for each handler + context_types = ContextTypes(context=CustomCallbackContext) + application = ( + Application.builder() + .token(token) + .context_types(context_types) + .defaults(defaults) + .build() + ) + application.bot_data["storage"] = storage + register_handlers(application) return application @@ -562,6 +490,7 @@ def register_handlers(application: Application): CommandHandler("changeplaylistimage", change_playlist_image), CommandHandler("playlistlink", send_playlist_link), CommandHandler("unlink", unlink_credentials), + CommandHandler("debug", debug_stuff), MessageHandler( filters.TEXT & filters.Regex(spotify_link_pattern), handle_spotify_links ), diff --git a/custom_callback_contex.py b/custom_callback_contex.py new file mode 100644 index 0000000..7412342 --- /dev/null +++ b/custom_callback_contex.py @@ -0,0 +1,11 @@ +from telegram.ext import CallbackContext as BaseCallbackContext + +from storage.dynamodb_storage import DynamoDBStorage + + +class CustomCallbackContext(BaseCallbackContext): + def __init__(self, application, chat_id=None, user_id=None): + super().__init__(application=application, chat_id=chat_id, user_id=user_id) + + def storage(self) -> DynamoDBStorage: + return self.bot_data["storage"] diff --git a/deploy_lambda.sh b/deploy_lambda.sh index 7d09481..ad6ce44 100755 --- a/deploy_lambda.sh +++ b/deploy_lambda.sh @@ -23,7 +23,7 @@ then fi rm -f $ZIP_FILE -zip -r $ZIP_FILE *.py html/ +zip -r $ZIP_FILE *.py html/ storage/ # Append the contents of the 'dependencies' directory at the root of the zip file if [ -d "dependencies" ]; then diff --git a/lambda_main.py b/lambda_main.py index 80fc795..c08bd21 100644 --- a/lambda_main.py +++ b/lambda_main.py @@ -7,6 +7,7 @@ import json import traceback from telegram import Update +from storage.dynamodb_storage import DynamoDBStorage if logging.getLogger().hasHandlers(): logging.getLogger().setLevel(logging.INFO) @@ -14,6 +15,7 @@ logging.basicConfig(level=logging.INFO) logger = logging.getLogger() +storage = DynamoDBStorage() def lambda_handler(event, context): logger.info(f"Event body type: {type(event)}") @@ -36,7 +38,7 @@ def handle_spotify_event(event): if not state_encoded or not code: return {"statusCode": 400, "body": "Missing required parameters"} - handle_spotify_auth(state_encoded, code) + handle_spotify_auth(state_encoded, code, storage) html_content = load_html_file("index.html") return { "statusCode": 200, @@ -47,7 +49,8 @@ def handle_spotify_event(event): async def main(event, context): TOKEN = os.getenv("TELEGRAM_BOT_TOKEN") - application = build_application(TOKEN) + + application = build_application(TOKEN, storage) # Convert the incoming event to a Telegram Update object if isinstance(event["body"], str): diff --git a/polling_main.py b/polling_main.py index e36a74d..8bdd052 100644 --- a/polling_main.py +++ b/polling_main.py @@ -5,6 +5,8 @@ from bot import handle_spotify_auth import logging from flask import Response +from storage.dynamodb_storage import DynamoDBStorage + webserver = Flask(__name__) @@ -24,7 +26,7 @@ def callback(): code = request.args.get("code") state = request.args.get("state") logger.info(f"honey we got the code: {code}, and the state: {state}") - handle_spotify_auth(state, code) + handle_spotify_auth(state, code, storage) return Response(html_content, status=200, content_type="text/html") @@ -34,17 +36,14 @@ def run_flask_app(): if __name__ == "__main__": TOKEN = os.getenv("TELEGRAM_BOT_TOKEN") - application = build_application(TOKEN) + + # Initialize storage + storage = DynamoDBStorage() + + application = build_application(TOKEN, storage) # Start the Flask app in a separate thread threading.Thread(target=run_flask_app, daemon=True).start() # Start polling application.run_polling() - # Main thread does something totally different. - # i = 0 - # while True: - # time.sleep(1) - # i += 1 - # sys.stdout.write(f"{i} {SPINNER[i % len(SPINNER)]}\r") - # sys.stdout.flush() diff --git a/requirements.txt b/requirements.txt index 564564f..568a6d7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,13 +10,13 @@ docutils==0.16 exceptiongroup==1.2.0 h11==0.14.0 httpcore==1.0.2 -httpx==0.25.2 +httpx~=0.26.0 idna==3.6 jmespath==1.0.1 pyasn1==0.5.1 python-dateutil==2.8.2 python-dotenv==1.0.0 -python-telegram-bot==20.7 +python-telegram-bot>=20.8 PyYAML==6.0.1 redis==5.0.1 requests==2.31.0 diff --git a/storage/dynamodb_storage.py b/storage/dynamodb_storage.py new file mode 100644 index 0000000..e710145 --- /dev/null +++ b/storage/dynamodb_storage.py @@ -0,0 +1,158 @@ +import boto3 +import logging +from .storage_interface import BotState, Storage +import os + + +class DynamoDBStorage(Storage): + def __init__(self): + dynamodb = boto3.resource("dynamodb", region_name="us-east-1") + self.bot_table = dynamodb.Table(os.getenv("BOT_TABLE")) + self.credentials_table = dynamodb.Table(os.getenv("CREDENTIALS_TABLE")) + + def save_current_state(self, chat_id, state_key: BotState): + try: + user_id = self.get_user_id_from_chat_id(chat_id) + if user_id is None: + user_id = self.get_user_id_from_channel_credentials(chat_id) + + if state_key is BotState.NO_STATE: + response = self.bot_table.update_item( + Key={"chat_id": str(chat_id)}, + UpdateExpression="REMOVE current_state", + ReturnValues="UPDATED_NEW", + ) + else: + response = self.bot_table.update_item( + Key={"chat_id": str(chat_id)}, + UpdateExpression="SET current_state = :state, user_id = :uid", + ExpressionAttributeValues={ + ":state": state_key.value, + ":uid": user_id, + }, + ReturnValues="UPDATED_NEW", + ) + logging.info(f"Updated state in DynamoDB: {response}") + return response + except Exception as e: + logging.error(f"Error saving current state to DynamoDB: {e}") + + def get_current_state(self, chat_id): + try: + response = self.bot_table.get_item(Key={"chat_id": str(chat_id)}) + if "Item" in response and "current_state" in response["Item"]: + state_value = response["Item"]["current_state"] + return BotState(state_value) + return None + except Exception as e: + logging.error(f"Error retrieving current state from DynamoDB: {e}") + return None + + def save_playlist_to_dynamodb(self, chat_id, playlist_id): + try: + self.bot_table.put_item( + Item={ + "chat_id": str(chat_id), + "playlist_id": playlist_id, + "user_id": self.get_user_id_from_chat_id(chat_id), + } + ) + except Exception as e: + logging.error(f"Error saving to DynamoDB: {e}") + + def get_playlist_from_dynamodb(self, chat_id): + try: + response = self.bot_table.get_item(Key={"chat_id": str(chat_id)}) + if "Item" in response: + playlist_id = response["Item"].get("playlist_id") + return playlist_id + return None + except Exception as e: + logging.error(f"Error retrieving from DynamoDB: {e}") + return None + + def get_user_id_from_chat_id(self, chat_id): + try: + response = self.bot_table.get_item(Key={"chat_id": str(chat_id)}) + if "Item" in response and "user_id" in response["Item"]: + return response["Item"]["user_id"] + else: + logging.info( + f"get_user_id_from_chat_id: No user_id found for chat_id: {chat_id}: {response}" + ) + return None + except Exception as e: + logging.error( + f"Error retrieving user ID from DynamoDB for chat_id: {chat_id}, error: {e}" + ) + return None + + def get_user_id_from_channel_credentials(self, chat_id): + try: + response = self.credentials_table.get_item(Key={"chat_id": str(chat_id)}) + if "Item" in response and "user_id" in response["Item"]: + return response["Item"]["user_id"] + else: + logging.info( + f"get_user_id_from_channel_credentials: No user_id found for chat_id: {chat_id}: {response}" + ) + return None + except Exception as e: + logging.error( + f"Error retrieving user ID from DynamoDB for chat_id: {chat_id}, error: {e}" + ) + return None + + def get_cached_token(self, chat_id): + try: + response = self.credentials_table.get_item(Key={"chat_id": str(chat_id)}) + if "Item" in response: + return response["Item"] + return None + except Exception as e: + logging.error(f"Error retrieving cached token from DynamoDB: {e}") + raise + + def save_token_to_cache(self, chat_id, user_id, token_info): + try: + self.credentials_table.put_item( + Item={ + "chat_id": str(chat_id), + "user_id": user_id, + **token_info, + } + ) + self.bot_table.update_item( + Key={"chat_id": str(chat_id)}, + UpdateExpression="SET user_id = :uid", + ExpressionAttributeValues={":uid": str(user_id)}, + ) + except Exception as e: + logging.error(f"Error saving token to DynamoDB: {e}") + raise + + def delete_item(self, table, chat_id): + try: + if table == "credentials_table": + response = self.credentials_table.delete_item(Key={"chat_id": str(chat_id)}) + logging.info(f"Deleted item from credentials_table for chat_id {chat_id}: {response}") + elif table == "bot_table": + response = self.bot_table.delete_item(Key={"chat_id": str(chat_id)}) + logging.info(f"Deleted item from bot_table for chat_id {chat_id}: {response}") + except Exception as e: + logging.error(f"Error deleting item from DynamoDB: {e}") + + def check_item_exists(self, table, chat_id): + try: + if table == "credentials_table": + response = self.credentials_table.get_item(Key={"chat_id": str(chat_id)}) + logging.info(f"Checked item in credentials_table for chat_id {chat_id}: {response}") + return "Item" in response + elif table == "bot_table": + response = self.bot_table.get_item(Key={"chat_id": str(chat_id)}) + logging.info(f"Checked item in bot_table for chat_id {chat_id}: {response}") + return "Item" in response + return False + except Exception as e: + logging.error(f"Error checking item in DynamoDB: {e}") + return False diff --git a/storage/in_memory_storage.py b/storage/in_memory_storage.py new file mode 100644 index 0000000..22d5091 --- /dev/null +++ b/storage/in_memory_storage.py @@ -0,0 +1,116 @@ +import logging +from .storage_interface import BotState, Storage + + +class InMemoryStorage(Storage): + def __init__(self): + self.bot_table = {} + self.credentials_table = {} + + def save_current_state(self, chat_id, state_key: BotState): + try: + user_id = self.get_user_id_from_chat_id(chat_id) + if user_id is None: + user_id = self.get_user_id_from_channel_credentials(chat_id) + + if state_key is BotState.NO_STATE: + if str(chat_id) in self.bot_table: + self.bot_table[str(chat_id)].pop("current_state", None) + else: + if str(chat_id) not in self.bot_table: + self.bot_table[str(chat_id)] = {} + self.bot_table[str(chat_id)]["current_state"] = state_key.value + self.bot_table[str(chat_id)]["user_id"] = user_id + + logging.info( + f"Updated state in InMemoryStorage: {self.bot_table[str(chat_id)]}" + ) + except Exception as e: + logging.error(f"Error saving current state to InMemoryStorage: {e}") + + def get_current_state(self, chat_id): + try: + item = self.bot_table.get(str(chat_id), {}) + if "current_state" in item: + state_value = item["current_state"] + return BotState(state_value) + return None + except Exception as e: + logging.error(f"Error retrieving current state from InMemoryStorage: {e}") + return None + + def save_playlist_to_dynamodb(self, chat_id, playlist_id): + try: + if str(chat_id) not in self.bot_table: + self.bot_table[str(chat_id)] = {} + self.bot_table[str(chat_id)]["playlist_id"] = playlist_id + self.bot_table[str(chat_id)]["user_id"] = self.get_user_id_from_chat_id( + chat_id + ) + except Exception as e: + logging.error(f"Error saving to InMemoryStorage: {e}") + + def get_playlist_from_dynamodb(self, chat_id): + try: + item = self.bot_table.get(str(chat_id), {}) + return item.get("playlist_id") + except Exception as e: + logging.error(f"Error retrieving from InMemoryStorage: {e}") + return None + + def get_user_id_from_chat_id(self, chat_id): + try: + item = self.bot_table.get(str(chat_id), {}) + return item.get("user_id") + except Exception as e: + logging.error( + f"Error retrieving user ID from InMemoryStorage for chat_id: {chat_id}, error: {e}" + ) + return None + + def get_user_id_from_channel_credentials(self, chat_id): + try: + item = self.credentials_table.get(str(chat_id), {}) + return item.get("user_id") + except Exception as e: + logging.error( + f"Error retrieving user ID from InMemoryStorage for chat_id: {chat_id}, error: {e}" + ) + return None + + def get_cached_token(self, chat_id): + try: + return self.credentials_table.get(str(chat_id)) + except Exception as e: + logging.error(f"Error retrieving cached token from InMemoryStorage: {e}") + raise + + def save_token_to_cache(self, chat_id, user_id, token_info): + try: + self.credentials_table[str(chat_id)] = { + "chat_id": str(chat_id), + "user_id": user_id, + **token_info, + } + if str(chat_id) not in self.bot_table: + self.bot_table[str(chat_id)] = {} + self.bot_table[str(chat_id)]["user_id"] = str(user_id) + except Exception as e: + logging.error(f"Error saving token to InMemoryStorage: {e}") + raise + + def delete_item(self, table, chat_id): + try: + table.pop(str(chat_id), None) + except Exception as e: + logging.error(f"Error deleting item from InMemoryStorage: {e}") + raise + + def check_item_exists(self, table, chat_id): + try: + exists = str(chat_id) in table + logging.info(f"check_item_exists: {exists}") + return exists + except Exception as e: + logging.error(f"Error checking item in InMemoryStorage: {e}") + raise diff --git a/storage/storage_interface.py b/storage/storage_interface.py new file mode 100644 index 0000000..9f998f9 --- /dev/null +++ b/storage/storage_interface.py @@ -0,0 +1,36 @@ +from abc import ABC, abstractmethod +from enum import Enum + + +class BotState(Enum): + AWAITING_PLAYLIST_IMAGE = "awaiting_playlist_image" + CHANGING_PLAYLIST_NAME = "changing_playlist_name" + CREATING_PLAYLIST = "creating_playlist" + CHANGING_PLAYLIST_IMAGE = "changing_playlist_image" + NO_STATE = None + + +class Storage(ABC): + @abstractmethod + def save_current_state(self, chat_id, state_key): + pass + + @abstractmethod + def get_current_state(self, chat_id): + pass + + @abstractmethod + def save_playlist_to_dynamodb(self, chat_id, playlist_id): + pass + + @abstractmethod + def get_playlist_from_dynamodb(self, chat_id): + pass + + @abstractmethod + def get_user_id_from_chat_id(self, chat_id): + pass + + @abstractmethod + def get_user_id_from_channel_credentials(self, chat_id): + pass diff --git a/test_telegram_bot.py b/test_telegram_bot.py index 9de3bbb..2eebfcd 100644 --- a/test_telegram_bot.py +++ b/test_telegram_bot.py @@ -1,63 +1,221 @@ +import json import unittest +from datetime import datetime +from typing import Optional, Tuple from unittest.mock import AsyncMock, MagicMock -from telegram import Update, Message, User, Chat -from telegram.ext import CallbackContext -from bot import start, help_command +from telegram import Update, Message, User, Chat, MessageEntity +from telegram._utils.defaultvalue import DEFAULT_NONE +from telegram._utils.types import ODVInput +from telegram.ext import CallbackContext, Application, ContextTypes +from telegram.request import BaseRequest, RequestData -class TestStartCommand(unittest.IsolatedAsyncioTestCase): - async def test_start_command(self): - update = MagicMock(spec=Update) +import bot +from bot import start, help_command, register_handlers +from custom_callback_contex import CustomCallbackContext +from storage.in_memory_storage import InMemoryStorage - update.message = MagicMock(spec=Message) - update.message.chat_id = 12345 - update.message.reply_text = AsyncMock() - context = MagicMock(spec=CallbackContext) +def counter(initial=1): + """counter creates a Python generator yielding incrementing integers.""" + i = initial + while True: + i += 1 + yield i - await start(update, context) - update.message.reply_text.assert_awaited_once_with( - "Hiya! I'm your Spotify Skunk bot 🦨. /createplaylist to add songs to your playlist!" +class UnexpectedTelegramEndpoint(Exception): + pass + + +class UnexpectedTelegramApiCall(Exception): + pass + + +class FakeTelegramApiServer(BaseRequest): + """FakeTelegramApiServerRequest implements a fake Telegram API server by overriding the default networking + implementation with one that operates entirely in-memory. + You can extract the outbound requests by inspecting the .messages_sent list. + """ + + fake_api_endpoint = "https://localhost/" + fake_token = "faketoken" + + def __init__(self, bot_user: User): + """Initialize a FakeTelegramApiServer. + :arg bot_user Describes the identity of the bot, as would be represented by the Telegram API server. + """ + super().__init__() + self.bot_user = bot_user + # Collected outbound requests will be stored here. + self.messages_sent = [] + + async def initialize(self) -> None: + pass + + async def shutdown(self) -> None: + pass + + async def do_request( + self, + url: str, + method: str, + request_data: Optional[RequestData] = None, + read_timeout: ODVInput[float] = DEFAULT_NONE, + write_timeout: ODVInput[float] = DEFAULT_NONE, + connect_timeout: ODVInput[float] = DEFAULT_NONE, + pool_timeout: ODVInput[float] = DEFAULT_NONE, + ) -> Tuple[int, bytes]: + """Implements a fake networking request/response.""" + expected_prefix = self.fake_api_endpoint + self.fake_token + "/" + if not url.startswith(expected_prefix): + raise UnexpectedTelegramEndpoint( + f"Expected prefix of {expected_prefix} but got URL {url}" + ) + path = url[len(expected_prefix) :] + + # AFAICT the library is only sending sendMessage and getMe requests, so these are the only two backend + # methods we fake. + if path == "getMe": + return 200, json.dumps( + {"ok": True, "result": json.loads(self.bot_user.to_json())} + ).encode("utf-8") + elif path == "sendMessage": + print(f"Sending message: {request_data.parameters}") + self.messages_sent.extend([request_data.parameters]) + # Fake a response from the Telegram servers. Note: This is not the full response that it sends to a + # message. The real response would contain the full message that was sent. + return 200, json.dumps({"ok": True, "result": {}}).encode("utf-8") + + # If the bot tries to make a Telegram API call to a different endpoint, just raise an exception and fail the + # test. + raise UnexpectedTelegramApiCall( + f"_request_wrapper(url={path}, method={method}, request_data={request_data.json_payload}" ) -class TestHelpCommand(unittest.IsolatedAsyncioTestCase): - async def setUp(self): - chat = Chat(id=12345, type="public") - user = User(id=67890, is_bot=False, first_name="Test User") - message = Message(message_id=1, date=1609459200, chat=chat, from_user=user) - self.update = Update(update_id=1, message=message) - - async def test_help_command(self): - context = CallbackContext.from_update(self.update, bot=None) - context = CallbackContext.from_update(self.update, bot=None) - context.bot = MagicMock() - context.bot_data = {} - context.user_data = {} - context.chat_data = {} - context.match = None - - self.update.message.reply_text = AsyncMock() - - # Invoke the help command - await help_command(self.update, context) - - # Check the response - expected_text = ( - "Here are the commands you can use:\n" - "/start - Start interacting with the bot\n" - "/createplaylist - Create a new Spotify playlist\n" - "/changeplaylistname - Change the name of the current playlist\n" - "/changeplaylistimage - Change the image of the current playlist\n" - "/resetplaylist - Reset so you can create a new playlist\n" - "/unlink - Unlink your Spotify credentials\n" - "/playlistlink - Get the link to the current playlist\n" - "/help - Show this help message\n" - "\nJust send me a Spotify track link to add it to your playlist!" +class TestBotWithoutMocks(unittest.IsolatedAsyncioTestCase): + """TestBotWithoutMocks allows for testing the responses generated when sending commands to the bot.""" + + # Utility values to generate incrementing values for the messages we generate. + __message_id_counter = counter() + __update_id_counter = counter(10000) + + async def asyncSetUp(self): + super().setUp() + # Note: it would be nice to test group chats AND 1:1 chats but for now we are just testing group chats. + self.chat = Chat(id=-1, type="public") + self.bot_user = User( + id=999, is_bot=False, first_name="Test", username="testskunkbot" + ) + self.request_collector = FakeTelegramApiServer(self.bot_user) + + # This must match the behavior in bot.build_application /except/ for the fields that are commented. + # TODO: unify with build_application. + self.app = ( + Application.builder() + .base_url( + self.request_collector.fake_api_endpoint + ) # disables interaction with telegram servers + .context_types(ContextTypes(context=CustomCallbackContext)) + .defaults(bot.defaults) + .get_updates_request( + self.request_collector + ) # replaces network I/O with FakeTelegramApiServerRequest + .request( + self.request_collector + ) # replaces network I/O with FakeTelegramApiServerRequest + .token(self.request_collector.fake_token) + .updater(None) # disables the library's polling or webhook behaviors + .build() ) - self.update.message.reply_text.assert_awaited_once_with(expected_text) + # Use an in-memory database rather than dynamodb + self.storage = InMemoryStorage() + self.app.bot_data["storage"] = self.storage + + register_handlers(self.app) + + await self.app.initialize() + + async def asyncTearDown(self): + await self.app.shutdown() + + async def sendMessage(self, text): + """sendMessage sends a message to the bot. + If the text starts with /, we assume it is a "Command" and send a slightly more special type of message to + deal with that. Otherwise, it is just a plain text message. + Returns the last message we received from the bot, or None if no messages were generated by the bot. + """ + + # Bot commands need to be annotated with entities=. + if text.startswith("/"): + message = Message( + chat=self.chat, + date=datetime(year=2024, month=1, day=1), + entities=[ + MessageEntity(type="bot_command", offset=0, length=len(text)) + ], + from_user=self.bot_user, + message_id=next(self.__message_id_counter), + text=text, + ) + else: + # Plain text commands do not need entities=. + message = Message( + chat=self.chat, + date=datetime(year=2024, month=1, day=1), + from_user=self.bot_user, + message_id=next(self.__message_id_counter), + text=text, + ) + + # `update` is the Update message that we would expect to receive from the Telegram API servers. + update = Update(update_id=next(self.__update_id_counter), message=message) + + # This app.process_update will trigger the appropriate command based on the user's message. Once it is + # complete, we can inspect the request_collector to determine which API calls were sent to the Telegram API + # servers. + await self.app.process_update( + Update.de_json(data=update.to_dict(), bot=self.app.bot) + ) + + if not self.request_collector.messages_sent: + return None + return self.request_collector.messages_sent.pop() + + async def test_help(self): + response = await self.sendMessage("/help") + self.assertTrue(response["text"].startswith("Here are the commands")) + + async def test_irrelevant_message(self): + response = await self.sendMessage("bark bark bark") + # This message is ignored, so there is no response. + self.assertIsNone(response) + + async def test_unlink_before_link(self): + response = await self.sendMessage("/unlink") + # TODO: Fix this behavior -- /unlink on an unlinked channel should not reply with something that looks like + # an error message. + self.assertEqual( + "Failed to unlink your Spotify credentials: credentials not found.", + response["text"], + ) + # Confirm that there are no credentials for this channel. + self.assertIsNone( + self.storage.get_user_id_from_channel_credentials(self.chat.id) + ) + # Confirm that there is no state for this channel. + self.assertIsNone(self.storage.get_current_state(self.chat.id)) + + async def test_playlistlink(self): + response = await self.sendMessage("/playlistlink") + self.assertEqual("No playlist found for this chat.", response["text"]) + + async def test_link(self): + response = await self.sendMessage("/createplaylist") + # TODO: verify the response contains a spotify URL + # self.assertTrue(response["text"] contains "http...") if __name__ == "__main__": unittest.main()