451 lines
16 KiB
Python
451 lines
16 KiB
Python
# test_utils.py
|
|
|
|
import unittest
|
|
import os
|
|
import re
|
|
import shutil
|
|
import tempfile
|
|
from unittest.mock import Mock, AsyncMock, patch
|
|
from datetime import datetime
|
|
|
|
# Adjusted import statement
|
|
from telegram_downloader_bot.utils import (
|
|
sanitize_name,
|
|
get_user_folder,
|
|
handle_media_message_contents,
|
|
download_tt_video,
|
|
make_fs,
|
|
extract_urls,
|
|
filter_tt_urls,
|
|
)
|
|
|
|
from pyrogram.types import Message, User, Chat
|
|
from pyrogram import Client
|
|
|
|
|
|
class TestSanitizeName(unittest.TestCase):
|
|
def test_alphanumeric_input(self):
|
|
input_str = "JohnDoe123"
|
|
expected_output = "JohnDoe123"
|
|
self.assertEqual(sanitize_name(input_str), expected_output)
|
|
|
|
def test_input_with_special_chars(self):
|
|
input_str = "John Doe!@#"
|
|
expected_output = "John_Doe"
|
|
self.assertEqual(sanitize_name(input_str), expected_output)
|
|
|
|
def test_input_with_only_special_chars(self):
|
|
input_str = "!@#$%^&*()"
|
|
expected_output = ""
|
|
self.assertEqual(sanitize_name(input_str), expected_output)
|
|
|
|
def test_empty_input(self):
|
|
input_str = ""
|
|
expected_output = ""
|
|
self.assertEqual(sanitize_name(input_str), expected_output)
|
|
|
|
|
|
class TestGetUserFolder(unittest.TestCase):
|
|
def setUp(self):
|
|
# Create a temporary directory for each test
|
|
self.tmp_path = tempfile.mkdtemp()
|
|
self.settings_patcher = patch('telegram_downloader_bot.settings.settings.storage', self.tmp_path)
|
|
self.settings_patcher.start()
|
|
|
|
def tearDown(self):
|
|
self.settings_patcher.stop()
|
|
# Remove the directory after the test
|
|
shutil.rmtree(self.tmp_path)
|
|
|
|
def test_forward_from_full_name(self):
|
|
user = Mock()
|
|
user.first_name = "John"
|
|
user.last_name = "Doe"
|
|
user.id = 12345
|
|
|
|
message = Mock()
|
|
message.forward_from = user
|
|
message.forward_from_chat = None
|
|
message.from_user = None
|
|
|
|
result = get_user_folder(message)
|
|
expected_folder = os.path.join(self.tmp_path, "telegram", "John_Doe")
|
|
self.assertEqual(result, expected_folder)
|
|
self.assertTrue(os.path.exists(expected_folder))
|
|
|
|
def test_forward_from_first_name_only(self):
|
|
user = Mock()
|
|
user.first_name = "John"
|
|
user.last_name = None
|
|
user.id = 12345
|
|
|
|
message = Mock()
|
|
message.forward_from = user
|
|
message.forward_from_chat = None
|
|
message.from_user = None
|
|
|
|
result = get_user_folder(message)
|
|
expected_folder = os.path.join(self.tmp_path, "telegram", "12345")
|
|
self.assertEqual(result, expected_folder)
|
|
self.assertTrue(os.path.exists(expected_folder))
|
|
|
|
def test_forward_from_chat_title(self):
|
|
chat = Mock()
|
|
chat.title = "My *Awesome* Group/Chat!"
|
|
|
|
message = Mock()
|
|
message.forward_from = None
|
|
message.forward_from_chat = chat
|
|
message.from_user = None
|
|
|
|
result = get_user_folder(message)
|
|
expected_folder = os.path.join(
|
|
self.tmp_path, "telegram", "My_Awesome_GroupChat"
|
|
)
|
|
self.assertEqual(result, expected_folder)
|
|
self.assertTrue(os.path.exists(expected_folder))
|
|
|
|
def test_from_user_full_name(self):
|
|
user = Mock()
|
|
user.first_name = "Jane"
|
|
user.last_name = "Doe"
|
|
user.id = 54321
|
|
|
|
message = Mock()
|
|
message.forward_from = None
|
|
message.forward_from_chat = None
|
|
message.from_user = user
|
|
|
|
result = get_user_folder(message)
|
|
expected_folder = os.path.join(self.tmp_path, "telegram", "Jane_Doe")
|
|
self.assertEqual(result, expected_folder)
|
|
self.assertTrue(os.path.exists(expected_folder))
|
|
|
|
def test_from_user_id(self):
|
|
user = Mock()
|
|
user.first_name = None
|
|
user.last_name = None
|
|
user.id = 54321
|
|
|
|
message = Mock()
|
|
message.forward_from = None
|
|
message.forward_from_chat = None
|
|
message.from_user = user
|
|
|
|
result = get_user_folder(message)
|
|
expected_folder = os.path.join(self.tmp_path, "telegram", "54321")
|
|
self.assertEqual(result, expected_folder)
|
|
self.assertTrue(os.path.exists(expected_folder))
|
|
|
|
|
|
# class TestHandleMediaMessageContents(unittest.IsolatedAsyncioTestCase):
|
|
# def setUp(self):
|
|
# # Create a temporary directory for each test
|
|
# self.tmp_path = tempfile.mkdtemp()
|
|
# self.settings_patcher = patch('telegram_downloader_bot.settings.settings.storage', self.tmp_path)
|
|
# self.settings_patcher.start()
|
|
|
|
# def tearDown(self):
|
|
# # Stop patching settings.storage
|
|
# self.settings_patcher.stop()
|
|
# # Remove the directory after the test
|
|
# shutil.rmtree(self.tmp_path)
|
|
|
|
# @patch('telegram_downloader_bot.utils.get_user_folder')
|
|
# async def test_handle_video(self, mock_get_user_folder):
|
|
# user_folder = os.path.join(self.tmp_path, "user_folder")
|
|
# mock_get_user_folder.return_value = user_folder
|
|
# os.makedirs(user_folder, exist_ok=True)
|
|
|
|
# client = Mock(spec=Client)
|
|
# client.download_media = AsyncMock()
|
|
|
|
# message = Mock(spec=Message)
|
|
# message.document = None
|
|
# message.photo = None
|
|
# message.video = Mock()
|
|
# message.video.file_id = "video_file_id"
|
|
# message.animation = None
|
|
# message.reply_text = AsyncMock()
|
|
|
|
# await handle_media_message_contents(client, message)
|
|
|
|
# expected_file_name = f"video_{message.video.file_id}.mp4"
|
|
# expected_file_path = os.path.join(user_folder, expected_file_name)
|
|
# client.download_media.assert_awaited_once_with(
|
|
# message, expected_file_path)
|
|
# message.reply_text.assert_awaited_once_with(
|
|
# f"Video saved to {user_folder}")
|
|
|
|
# @patch('telegram_downloader_bot.utils.get_user_folder')
|
|
# async def test_handle_animation(self, mock_get_user_folder):
|
|
# user_folder = os.path.join(self.tmp_path, "user_folder")
|
|
# mock_get_user_folder.return_value = user_folder
|
|
# os.makedirs(user_folder, exist_ok=True)
|
|
|
|
# client = Mock(spec=Client)
|
|
# client.download_media = AsyncMock()
|
|
|
|
# message = Mock(spec=Message)
|
|
# message.document = None
|
|
# message.photo = None
|
|
# message.video = None
|
|
# message.animation = Mock()
|
|
# message.animation.file_id = "animation_file_id"
|
|
# message.reply_text = AsyncMock()
|
|
|
|
# await handle_media_message_contents(client, message)
|
|
|
|
# expected_file_name = f"gif_{message.animation.file_id}.gif"
|
|
# expected_file_path = os.path.join(user_folder, expected_file_name)
|
|
# client.download_media.assert_awaited_once_with(
|
|
# message.animation, expected_file_path)
|
|
# message.reply_text.assert_awaited_once_with(
|
|
# f"GIF saved to {user_folder}")
|
|
|
|
# @patch('telegram_downloader_bot.utils.get_user_folder')
|
|
# async def test_handle_document(self, mock_get_user_folder):
|
|
# user_folder = os.path.join(self.tmp_path, "user_folder")
|
|
# mock_get_user_folder.return_value = user_folder
|
|
# os.makedirs(user_folder, exist_ok=True)
|
|
|
|
# client = Mock(spec=Client)
|
|
# client.download_media = AsyncMock()
|
|
|
|
# message = Mock(spec=Message)
|
|
# message.document = Mock()
|
|
# message.document.file_name = "test_document.pdf"
|
|
# message.photo = None
|
|
# message.video = None
|
|
# message.animation = None
|
|
# message.reply_text = AsyncMock()
|
|
|
|
# await handle_media_message_contents(client, message)
|
|
|
|
# expected_file_path = os.path.join(user_folder, "test_document.pdf")
|
|
# client.download_media.assert_awaited_once_with(
|
|
# message, expected_file_path)
|
|
# message.reply_text.assert_awaited_once_with(
|
|
# f"Document saved to {user_folder}")
|
|
|
|
# @patch('telegram_downloader_bot.utils.get_user_folder')
|
|
# async def test_handle_photo(self, mock_get_user_folder):
|
|
# user_folder = os.path.join(self.tmp_path, "user_folder")
|
|
# mock_get_user_folder.return_value = user_folder
|
|
# os.makedirs(user_folder, exist_ok=True)
|
|
|
|
# client = Mock(spec=Client)
|
|
# client.download_media = AsyncMock()
|
|
|
|
# message = Mock(spec=Message)
|
|
# message.document = None
|
|
# message.photo = Mock()
|
|
# message.photo.file_id = "photo_file_id"
|
|
# message.video = None
|
|
# message.animation = None
|
|
# message.reply_text = AsyncMock()
|
|
|
|
# await handle_media_message_contents(client, message)
|
|
|
|
# expected_file_name = f"photo_{message.photo.file_id}.jpg"
|
|
# expected_file_path = os.path.join(user_folder, expected_file_name)
|
|
# client.download_media.assert_awaited_once_with(
|
|
# message.photo, expected_file_path)
|
|
# message.reply_text.assert_awaited_once_with(
|
|
# f"Photo saved to {user_folder}")
|
|
|
|
# @patch('telegram_downloader_bot.utils.get_user_folder')
|
|
# async def test_handle_unknown_media(self, mock_get_user_folder):
|
|
# user_folder = os.path.join(self.tmp_path, "user_folder")
|
|
# mock_get_user_folder.return_value = user_folder
|
|
# os.makedirs(user_folder, exist_ok=True)
|
|
|
|
# client = Mock(spec=Client)
|
|
# client.download_media = AsyncMock()
|
|
|
|
# message = Mock(spec=Message)
|
|
# message.document = None
|
|
# message.photo = None
|
|
# message.video = None
|
|
# message.animation = None
|
|
# message.reply_text = AsyncMock()
|
|
|
|
# await handle_media_message_contents(client, message)
|
|
|
|
# client.download_media.assert_not_called()
|
|
# message.reply_text.assert_awaited_once_with("Unknown media type!")
|
|
|
|
|
|
# class TestDownloadTTVideo(unittest.TestCase):
|
|
# def setUp(self):
|
|
# # Create a temporary directory for each test
|
|
# self.tmp_path = tempfile.mkdtemp()
|
|
# os.makedirs(os.path.join(self.tmp_path, "tiktok"), exist_ok=True)
|
|
# self.settings_patcher = patch("telegram_downloader_bot.settings.settings.storage", self.tmp_path)
|
|
# self.settings_patcher.start()
|
|
|
|
# # Paths to the valid and invalid video files
|
|
# self.valid_video_path = os.path.join(self.tmp_path, "valid.mp4")
|
|
# with open(self.valid_video_path, 'wb') as f:
|
|
# f.write(b'valid mp4 content')
|
|
|
|
# self.invalid_video_path = os.path.join(self.tmp_path, "invalid.mp4")
|
|
# with open(self.invalid_video_path, 'wb') as f:
|
|
# f.write(b'invalid mp4 content')
|
|
|
|
# def tearDown(self):
|
|
# self.settings_patcher.stop()
|
|
# # Remove the directory after the test
|
|
# shutil.rmtree(self.tmp_path)
|
|
|
|
# @patch('telegram_downloader_bot.utils.snaptik')
|
|
# @patch('telegram_downloader_bot.utils.datetime')
|
|
# def test_download_tt_video_with_valid_video(self, mock_datetime, mock_snaptik):
|
|
# # Mock datetime
|
|
# mock_now = datetime(2023, 1, 1, 12, 0, 0)
|
|
# mock_datetime.now.return_value = mock_now
|
|
|
|
# # Read the content of valid.mp4
|
|
# with open(self.valid_video_path, 'rb') as f:
|
|
# valid_video_content = f.read()
|
|
|
|
# # Mock snaptik to return a video that returns valid.mp4 content
|
|
# mock_video = Mock()
|
|
# mock_video.download.return_value.getbuffer.return_value = valid_video_content
|
|
# mock_snaptik.return_value = [mock_video]
|
|
|
|
# # Call the function
|
|
# download_tt_video("http://tiktok.com/video123")
|
|
|
|
# # Verify that the file was saved correctly
|
|
# video_filename = mock_now.strftime(
|
|
# "video-tiktok-%Y-%m-%d_%H-%M-%S.mp4")
|
|
# video_filepath = os.path.join(self.tmp_path, "tiktok", video_filename)
|
|
# self.assertTrue(os.path.exists(video_filepath))
|
|
|
|
# with open(video_filepath, 'rb') as f:
|
|
# content = f.read()
|
|
# self.assertEqual(content, valid_video_content)
|
|
|
|
# @patch('telegram_downloader_bot.utils.snaptik')
|
|
# @patch('telegram_downloader_bot.utils.datetime')
|
|
# def test_download_tt_video_with_invalid_video(self, mock_datetime, mock_snaptik):
|
|
# # Mock datetime
|
|
# mock_now = datetime(2023, 1, 1, 12, 0, 0)
|
|
# mock_datetime.now.return_value = mock_now
|
|
|
|
# # Read the content of invalid.mp4
|
|
# with open(self.invalid_video_path, 'rb') as f:
|
|
# invalid_video_content = f.read()
|
|
|
|
# # Mock snaptik to return a video that returns invalid.mp4 content
|
|
# mock_video = Mock()
|
|
# mock_video.download.return_value.getbuffer.return_value = invalid_video_content
|
|
# mock_snaptik.return_value = [mock_video]
|
|
|
|
# # Call the function
|
|
# download_tt_video("http://tiktok.com/video123")
|
|
|
|
# # Verify that the file was saved
|
|
# video_filename = mock_now.strftime(
|
|
# "video-tiktok-%Y-%m-%d_%H-%M-%S.mp4")
|
|
# video_filepath = os.path.join(self.tmp_path, "tiktok", video_filename)
|
|
# self.assertTrue(os.path.exists(video_filepath))
|
|
|
|
# with open(video_filepath, 'rb') as f:
|
|
# content = f.read()
|
|
# self.assertEqual(content, invalid_video_content)
|
|
|
|
# @patch('telegram_downloader_bot.utils.snaptik')
|
|
# @patch('telegram_downloader_bot.utils.datetime')
|
|
# def test_download_tt_video_no_videos(self, mock_datetime, mock_snaptik):
|
|
# # Mock datetime
|
|
# mock_now = datetime(2023, 1, 1, 12, 0, 0)
|
|
# mock_datetime.datetime.now.return_value = mock_now
|
|
|
|
# # Mock snaptik to return an empty list
|
|
# mock_snaptik.return_value = []
|
|
|
|
# # Call the function
|
|
# download_tt_video("http://tiktok.com/video123")
|
|
|
|
# # Verify that no files were created
|
|
# tiktok_folder = os.path.join(self.tmp_path, "tiktok")
|
|
# files = os.listdir(tiktok_folder)
|
|
# self.assertEqual(len(files), 0)
|
|
|
|
|
|
# class TestMakeFS(unittest.TestCase):
|
|
# def setUp(self):
|
|
# self.tmp_path = tempfile.mkdtemp()
|
|
|
|
# def tearDown(self):
|
|
# shutil.rmtree(self.tmp_path)
|
|
|
|
# def test_make_fs(self):
|
|
# make_fs(self.tmp_path)
|
|
# self.assertTrue(os.path.exists(os.path.join(self.tmp_path, "tiktok")))
|
|
# self.assertTrue(os.path.exists(
|
|
# os.path.join(self.tmp_path, "telegram")))
|
|
|
|
|
|
# class TestExtractURLs(unittest.TestCase):
|
|
# def test_no_urls(self):
|
|
# text = "This is some text without any URLs."
|
|
# result = extract_urls(text)
|
|
# self.assertEqual(result, [])
|
|
|
|
# def test_single_url(self):
|
|
# text = "Check out this link: http://example.com"
|
|
# result = extract_urls(text)
|
|
# self.assertEqual(result, ["http://example.com"])
|
|
|
|
# def test_multiple_urls(self):
|
|
# text = "Here are some links: http://example.com and https://test.com/page"
|
|
# result = extract_urls(text)
|
|
# self.assertEqual(
|
|
# result, ["http://example.com", "https://test.com/page"])
|
|
|
|
# def test_malformed_url(self):
|
|
# text = "This is not a URL: htt://badurl.com"
|
|
# result = extract_urls(text)
|
|
# self.assertEqual(result, [])
|
|
|
|
# def test_urls_with_special_chars(self):
|
|
# text = "Link: https://example.com/page?param=value#anchor"
|
|
# result = extract_urls(text)
|
|
# self.assertEqual(
|
|
# result, ["https://example.com/page?param=value#anchor"])
|
|
|
|
|
|
# class TestFilterTTURLs(unittest.TestCase):
|
|
# def test_empty_list(self):
|
|
# urls = []
|
|
# result = filter_tt_urls(urls)
|
|
# self.assertEqual(result, [])
|
|
|
|
# def test_no_tiktok_urls(self):
|
|
# urls = ["http://example.com", "https://test.com/page"]
|
|
# result = filter_tt_urls(urls)
|
|
# self.assertEqual(result, [])
|
|
|
|
# def test_mixed_urls(self):
|
|
# urls = [
|
|
# "http://example.com",
|
|
# "https://www.tiktok.com/@user/video/123",
|
|
# "http://tiktok.com/video1",
|
|
# "https://test.com/page",
|
|
# ]
|
|
# expected = [
|
|
# "https://www.tiktok.com/@user/video/123",
|
|
# "http://tiktok.com/video1",
|
|
# ]
|
|
# result = filter_tt_urls(urls)
|
|
# self.assertEqual(result, expected)
|
|
|
|
# def test_tiktok_in_query_params(self):
|
|
# urls = ["http://example.com?watch=tiktok", "https://other.com/path"]
|
|
# expected = ["http://example.com?watch=tiktok"]
|
|
# result = filter_tt_urls(urls)
|
|
# self.assertEqual(result, expected)
|