From 92d6e5d8bc8d67aafe4bcba0e6f2d187e0cd4a89 Mon Sep 17 00:00:00 2001 From: Sergei Vedishchev Date: Mon, 1 Apr 2024 23:59:53 +0200 Subject: [PATCH] fix bot action typing for non-related uploads --- engines/ideogram_img.py | 4 +- lambda/chatbot.py | 83 ++++++++++++++++++--------------------- tests/test_bing_images.py | 4 +- tests/test_ideogram.py | 8 ++++ 4 files changed, 52 insertions(+), 47 deletions(-) diff --git a/engines/ideogram_img.py b/engines/ideogram_img.py index 6151808..b0b97c4 100644 --- a/engines/ideogram_img.py +++ b/engines/ideogram_img.py @@ -50,9 +50,9 @@ def is_expired(id_token: str) -> bool: logging.info(f"exp:{exp}, now:{now}") return exp > now except jwt.ExpiredSignatureError: - return False + return True except jwt.InvalidTokenError: - return False + return True def refresh_iss_tokens(refresh_token: str) -> dict: diff --git a/lambda/chatbot.py b/lambda/chatbot.py index 2e61e76..17853a5 100644 --- a/lambda/chatbot.py +++ b/lambda/chatbot.py @@ -417,39 +417,49 @@ async def process_voice_message(update: Update, context: ContextTypes.DEFAULT_TY @send_typing_action +async def process_upload( + update: Update, context: ContextTypes.DEFAULT_TYPE, file_id: str, file_name: str +) -> None: + s3_bucket = read_ssm_param(param_name="BOT_S3_BUCKET") + file = await bot.get_file(file_id) + path = await upload_to_s3(file, s3_bucket, "att", file_name) + logging.info(f"File uploaded {path}") + user_id = int(update.effective_user.id) + config = user_config.read(user_id) + envelop = { + "type": "text", + "user_id": user_id, + "username": update.effective_user.name, + "update_id": update.update_id, + "message_id": update.effective_message.id, + "text": update.message.caption, + "chat_id": update.effective_chat.id, + "timestamp": update.effective_message.date.timestamp(), + "config": config, + "file": path, + } + # logging.info(envelop) + await __send_envelop(envelop, json.dumps(config["engines"])) + + async def process_photo(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: if update.message.photo is None: return logging.info("File upload in 'process_photo'") # logging.info(update.message) - caption = update.message.caption - if bot.name not in caption and "group" in update.message.chat.type: + if bot.name not in update.message.caption and "group" in update.message.chat.type: return - s3_bucket = read_ssm_param(param_name="BOT_S3_BUCKET") photo = max(update.message.photo, key=lambda x: x.file_size) logging.info(photo) file_id = photo.file_id logging.info(file_id) try: - file = await bot.get_file(file_id) - path = await upload_to_s3(file, s3_bucket, "att", f"{photo.file_unique_id}.jpg") - logging.info(f"File uploaded {path}") - user_id = int(update.effective_user.id) - config = user_config.read(user_id) - envelop = { - "type": "text", - "user_id": user_id, - "username": update.effective_user.name, - "update_id": update.update_id, - "message_id": update.effective_message.id, - "text": caption, - "chat_id": update.effective_chat.id, - "timestamp": update.effective_message.date.timestamp(), - "config": config, - "file": path, - } - # logging.info(envelop) - await __send_envelop(envelop, json.dumps(config["engines"])) + await process_upload( + update=update, + context=context, + file_id=file_id, + file_name=f"{photo.file_unique_id}.jpg", + ) except Exception as e: logging.error( msg="Exception occured during processing of the picture", @@ -462,34 +472,19 @@ async def process_attachment(update: Update, context: ContextTypes.DEFAULT_TYPE) if update.message is None: return logging.info(update.message) - caption = update.message.caption - if bot.name not in caption and "group" in update.message.chat.type: + if bot.name not in update.message.caption and "group" in update.message.chat.type: return attachment = update.message.effective_attachment logging.info(attachment) file_id = attachment.file_id logging.info(file_id) - s3_bucket = read_ssm_param(param_name="BOT_S3_BUCKET") try: - file = await bot.get_file(file_id) - path = await upload_to_s3(file, s3_bucket, "att", attachment.file_name) - logging.info(f"File uploaded {path}") - user_id = int(update.effective_user.id) - config = user_config.read(user_id) - envelop = { - "type": "text", - "user_id": user_id, - "username": update.effective_user.name, - "update_id": update.update_id, - "message_id": update.effective_message.id, - "text": caption, - "chat_id": update.effective_chat.id, - "timestamp": update.effective_message.date.timestamp(), - "config": config, - "file": path, - } - # logging.info(envelop) - await __send_envelop(envelop, json.dumps(config["engines"])) + await process_upload( + update=update, + context=context, + file_id=file_id, + file_name=attachment.file_name, + ) except Exception: logging.error( msg="Exception occured during processing of the attachment", diff --git a/tests/test_bing_images.py b/tests/test_bing_images.py index 312c915..befa9c3 100644 --- a/tests/test_bing_images.py +++ b/tests/test_bing_images.py @@ -1,7 +1,9 @@ +import pytest + from engines.dalle_img import imageGen -# @pytest.mark.skip() +@pytest.mark.skip() def test_check_and_refresh(capsys): with capsys.disabled(): prompt = "roasted coffee beans in palms" diff --git a/tests/test_ideogram.py b/tests/test_ideogram.py index 6d1a48b..6546f3f 100644 --- a/tests/test_ideogram.py +++ b/tests/test_ideogram.py @@ -43,3 +43,11 @@ def test_get_session_cookies(capsys): token = read_json_from_s3(bucket_name=bucket_name, file_name="google_auth.json") response = get_session_cookies(iss_token=token["access_token"]) assert response + + +@pytest.mark.skip() +def test_is_expired(capsys): + with capsys.disabled(): + bucket_name = read_ssm_param(param_name="BOT_S3_BUCKET") + token = read_json_from_s3(bucket_name=bucket_name, file_name="google_auth.json") + assert not is_expired(token["access_token"])