From 3f2ebe268d0287354512fdeb20e238d968f80137 Mon Sep 17 00:00:00 2001
From: Aleksandr Movchan <aleksandr@mobiuslabs.com>
Date: Fri, 9 Aug 2024 09:56:55 +0000
Subject: [PATCH] Update dependency for yt-dlp and added mock for youtube
 downloader.

---
 aana/tests/units/test_video.py | 107 +++++++++++++++++++++------------
 poetry.lock                    |   6 +-
 pyproject.toml                 |   2 +-
 3 files changed, 73 insertions(+), 42 deletions(-)

diff --git a/aana/tests/units/test_video.py b/aana/tests/units/test_video.py
index 12998b8c..c2e5f9b3 100644
--- a/aana/tests/units/test_video.py
+++ b/aana/tests/units/test_video.py
@@ -167,44 +167,6 @@ def test_download_video():
     finally:
         video.cleanup()
 
-    # Test Youtube URL
-    youtube_url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ"
-    youtube_url_hash = hashlib.md5(
-        youtube_url.encode(), usedforsecurity=False
-    ).hexdigest()
-    video_dir = settings.video_dir
-    expected_path = video_dir / f"{youtube_url_hash}.mp4"
-    # remove the file if it exists
-    expected_path.unlink(missing_ok=True)
-
-    try:
-        youtube_video_input = VideoInput(url=youtube_url, media_id="dQw4w9WgXcQ")
-        video = download_video(youtube_video_input)
-        assert isinstance(video, Video)
-        # ignore extension
-        assert video.path.with_suffix("") == expected_path.with_suffix("")
-        assert video.path is not None
-        assert video.path.exists()
-        assert video.content is None
-        assert video.url == youtube_url
-        assert video.media_id == "dQw4w9WgXcQ"
-        assert (
-            video.title
-            == "Rick Astley - Never Gonna Give You Up (Official Music Video)"
-        )
-        assert video.description.startswith(
-            "The official video for “Never Gonna Give You Up” by Rick Astley."
-        )
-    finally:
-        if video and video.path:
-            video.path.unlink(missing_ok=True)
-
-    # Test YoutubeVideoInput with invalid youtube_url
-    youtube_url = "https://www.youtube.com/watch?v=invalid_url"
-    youtube_video_input = VideoInput(url=youtube_url)
-    with pytest.raises(DownloadException):
-        download_video(youtube_video_input)
-
     # Test url that doesn't contain a video
     url = "https://mobius-public.s3.eu-west-1.amazonaws.com/Starry_Night.jpeg"
     video_input = VideoInput(url=url)
@@ -235,8 +197,20 @@ def test_download_video():
         ),
     ],
 )
-def test_get_video_metadata_success(url, title, description, duration):
+def test_get_video_metadata_success(mocker, url, title, description, duration):
     """Test getting video metadata."""
+    # Mock the yt-dlp YoutubeDL class and its methods
+    mock_ydl = mocker.patch("yt_dlp.YoutubeDL", autospec=True)
+    mock_ydl_instance = mock_ydl.return_value
+    mock_ydl_instance.__enter__.return_value = mock_ydl_instance
+
+    mock_info = {
+        "title": title,
+        "description": description,
+        "duration": duration,
+    }
+    mock_ydl_instance.extract_info.return_value = mock_info
+
     metadata = get_video_metadata(url)
     assert isinstance(metadata, VideoMetadata)
     assert metadata.title == title
@@ -249,3 +223,58 @@ def test_get_video_metadata_failure():
     url = "https://www.youtube.com/watch?v=invalid_url"
     with pytest.raises(DownloadException):
         get_video_metadata(url)
+
+
+def test_download_youtube_video(mocker):
+    """Test downloading a YouTube video."""
+    youtube_url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ"
+    youtube_url_hash = hashlib.md5(
+        youtube_url.encode(), usedforsecurity=False
+    ).hexdigest()
+    video_dir = settings.video_dir
+    expected_path = Path(video_dir) / f"{youtube_url_hash}.mp4"
+
+    # Mock the VideoInput object
+    youtube_video_input = VideoInput(url=youtube_url, media_id="dQw4w9WgXcQ")
+
+    # Mock the yt-dlp YoutubeDL class and its methods
+    mock_ydl = mocker.patch("yt_dlp.YoutubeDL", autospec=True)
+    mock_ydl_instance = mock_ydl.return_value
+    mock_ydl_instance.__enter__.return_value = mock_ydl_instance
+
+    mock_info = {
+        "title": "Rick Astley - Never Gonna Give You Up (Official Music Video)",
+        "description": "The official video for “Never Gonna Give You Up” by Rick Astley.",
+        "ext": "mp4",
+    }
+    mock_ydl_instance.extract_info.return_value = mock_info
+    mock_ydl_instance.prepare_filename.return_value = str(expected_path)
+    mock_ydl_instance.download.return_value = None
+
+    # # Mock Path.exists to simulate that the file will exist after download
+    mocker.patch("aana.integrations.external.yt_dlp.Path.exists", return_value=True)
+
+    # Mock decord.VideoReader to avoid trying to read the non-existent mocked file
+    mocker.patch("decord.VideoReader", autospec=True)
+
+    video = download_video(youtube_video_input)
+
+    assert isinstance(video, Video)
+    assert video.path.with_suffix("") == expected_path.with_suffix("")
+    assert video.path is not None
+    assert video.path.exists()
+    assert video.content is None
+    assert video.url == youtube_url
+    assert video.media_id == "dQw4w9WgXcQ"
+    assert video.title == "Rick Astley - Never Gonna Give You Up (Official Music Video)"
+    assert video.description.startswith(
+        "The official video for “Never Gonna Give You Up” by Rick Astley."
+    )
+
+
+def test_download_youtube_video_failure():
+    """Test YoutubeVideoInput with invalid youtube URL."""
+    youtube_url = "https://www.youtube.com/watch?v=invalid_url"
+    youtube_video_input = VideoInput(url=youtube_url)
+    with pytest.raises(DownloadException):
+        download_video(youtube_video_input)
diff --git a/poetry.lock b/poetry.lock
index c8def7e7..fd20d837 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.6.1 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
 
 [[package]]
 name = "accelerate"
@@ -4072,6 +4072,8 @@ files = [
     {file = "orjson-3.10.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:960db0e31c4e52fa0fc3ecbaea5b2d3b58f379e32a95ae6b0ebeaa25b93dfd34"},
     {file = "orjson-3.10.6-cp312-none-win32.whl", hash = "sha256:a6ea7afb5b30b2317e0bee03c8d34c8181bc5a36f2afd4d0952f378972c4efd5"},
     {file = "orjson-3.10.6-cp312-none-win_amd64.whl", hash = "sha256:874ce88264b7e655dde4aeaacdc8fd772a7962faadfb41abe63e2a4861abc3dc"},
+    {file = "orjson-3.10.6-cp313-none-win32.whl", hash = "sha256:efdf2c5cde290ae6b83095f03119bdc00303d7a03b42b16c54517baa3c4ca3d0"},
+    {file = "orjson-3.10.6-cp313-none-win_amd64.whl", hash = "sha256:8e190fe7888e2e4392f52cafb9626113ba135ef53aacc65cd13109eb9746c43e"},
     {file = "orjson-3.10.6-cp38-cp38-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:66680eae4c4e7fc193d91cfc1353ad6d01b4801ae9b5314f17e11ba55e934183"},
     {file = "orjson-3.10.6-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:caff75b425db5ef8e8f23af93c80f072f97b4fb3afd4af44482905c9f588da28"},
     {file = "orjson-3.10.6-cp38-cp38-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3722fddb821b6036fd2a3c814f6bd9b57a89dc6337b9924ecd614ebce3271394"},
@@ -8348,4 +8350,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.10,<4.0"
-content-hash = "42fd6e903fbef04cf1679137e2fba066f8b4ecf006d9b28e6f2cef313d1032b4"
+content-hash = "021a09ad1cd2d2202d9b31e3784e258f79cbda9ddf9d4b0e66f01f4558fc0968"
diff --git a/pyproject.toml b/pyproject.toml
index 92ff4df8..f5f0090f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -54,7 +54,7 @@ torch = ">=2.0.0"
 torchaudio = ">=2.0.0"
 torchvision = ">=0.15.0"
 vllm = "^0.5.0.post1"
-yt-dlp = ">=2024.4.9"
+yt-dlp = ">=2024.08.06"
 
 [tool.poetry.group.dev.dependencies]
 ipykernel = "^6.25.2"