- Add test for video_transcription command;

- Add some necessary improvements in other files

Author: aninhasalesp

tests/commands/test_channel_info.py

@@ -47,4 +47,6 @@ def test_channel_ids_from_urls_and_usernames(mocker, channels_urls, usernames):
     channel_id_from_username_mock.assert_has_calls(
         [call(username) for username in usernames]
     )
-    channels_infos_mock.assert_called_once_with([ids_from_urls_mock, ids_from_usernames_mock])
+    channels_infos_mock.assert_called_once()
+    assert ids_from_usernames_mock in channels_infos_mock.call_args.args[0]
+    assert ids_from_urls_mock in channels_infos_mock.call_args.args[0]

tests/commands/test_video_transcription.py

@@ -0,0 +1,56 @@
+from unittest.mock import Mock
+
+from youtool.commands import VideoTranscription
+
+
+def test_video_transcription(mocker, videos_ids, videos_urls, tmp_path):
+    youtube_mock = mocker.patch("youtool.commands.video_transcription.YouTube")
+
+    language_code = "pt_br"
+
+    videos_transcriptions_mock = Mock()
+    youtube_mock.return_value.videos_transcriptions = videos_transcriptions_mock
+
+    for video_id in videos_ids:
+        open(tmp_path / f"{video_id}.{language_code}.vtt", "a").close()
+
+    result = VideoTranscription.execute(
+        ids=videos_ids, urls=videos_urls, language_code=language_code, output_dir=tmp_path
+    )
+
+    videos_transcriptions_mock.assert_called_once_with(
+        list(set(videos_ids)), language_code, tmp_path
+    )
+
+    for video_id in videos_ids:
+        assert str(tmp_path / f"{video_id}.{language_code}.vtt") in result
+
+
+def test_video_transcription_input_from_file(mocker, videos_ids, tmp_path):
+    youtube_mock = mocker.patch("youtool.commands.video_transcription.YouTube")
+
+    language_code = "pt_br"
+
+    videos_transcriptions_mock = Mock()
+    youtube_mock.return_value.videos_transcriptions = videos_transcriptions_mock
+
+    input_file_path = tmp_path / "input_file.csv"
+
+    with open(input_file_path, "w") as input_csv:
+        input_csv.write("video_id\n" + "\n".join(videos_ids))
+
+    for video_id in videos_ids:
+        open(tmp_path / f"{video_id}.{language_code}.vtt", "a").close()
+
+    result = VideoTranscription.execute(
+        ids=None, urls=None,
+        language_code=language_code, output_dir=tmp_path,
+        input_file_path=input_file_path
+    )
+
+    videos_transcriptions_mock.assert_called_once_with(
+        list(set(videos_ids)), language_code, tmp_path
+    )
+
+    for video_id in videos_ids:
+        assert str(tmp_path / f"{video_id}.{language_code}.vtt") in result

tests/test_cli.py

@@ -1,5 +1,6 @@
 import pytest
 
+from pathlib import Path
 from subprocess import run
 
 from youtool.commands import COMMANDS
@@ -12,9 +13,13 @@
 )
 def test_missing_api_key(monkeypatch: pytest.MonkeyPatch, command: Command):
     monkeypatch.delenv('YOUTUBE_API_KEY', raising=False)
-    cli_path = "youtool/cli.py"
-    command = ["python", cli_path, command.name]
-    result = run(command, capture_output=True, text=True, check=False)
+    cli_path = Path("youtool") / "cli.py"
+    command_string = ["python", cli_path, command.name]
+    for arg in command.arguments:
+        if arg.get("required"):
+            command_string.append(arg.get("name"))
+            command_string.append("test_value")
+    result = run(command_string, capture_output=True, text=True, check=False)
 
     assert result.returncode == 2
-    assert "YouTube API Key is required" in result.stderr
+    assert "YouTube API Key is required" in result.stderr

youtool/commands/base.py

@@ -80,7 +80,11 @@ def execute(cls, **kwargs) -> str:  # noqa: D417
         raise NotImplementedError()
 
     @staticmethod
-    def data_from_csv(file_path: Path, data_column_name: Optional[str] = None) -> List[str]:
+    def data_from_csv(
+        file_path: Path,
+        data_column_name: Optional[str] = None,
+        raise_column_exception: bool = True
+    ) -> List[str]:
         """Extracts a list of URLs from a specified CSV file.
 
         Args:
@@ -107,7 +111,10 @@ def data_from_csv(file_path: Path, data_column_name: Optional[str] = None) -> Li
                 raise ValueError("Fieldnames is None")
 
             if data_column_name not in fieldnames:
-                raise Exception(f"Column {data_column_name} not found on {file_path}")
+                if raise_column_exception:
+                    raise Exception(f"Column {data_column_name} not found on {file_path}")
+                return data
+
             for row in reader:
                 value = row.get(data_column_name)
                 if value is not None:

youtool/commands/video_transcription.py

@@ -16,11 +16,12 @@ class VideoTranscription(Command):
         {"name": "--output-dir", "type": str, "help": "Output directory to save transcriptions"},
         {"name": "--language-code", "type": str, "help": "Language code for transcription"},
         {"name": "--api-key", "type": str, "help": "API key for YouTube Data API"},
+        {"name": "--url-column-name", "type": str, "help": "URL column name on csv input files"},
+        {"name": "--id-column-name", "type": str, "help": "Channel ID column name on csv output files"}
     ]
 
-    TRANSCRIPTION_COLUMNS: List[str] = [
-        "video_id", "transcription_text"
-    ]
+    ID_COLUMN_NAME: str = "video_id"
+    URL_COLUMN_NAME: str = "video_url"
 
     @classmethod
     def execute(cls, **kwargs) -> str:
@@ -38,33 +39,39 @@ def execute(cls, **kwargs) -> str:
         Returns:
             A message indicating the result of the command. Reports success or failure for each video transcription download.
         """
-        ids = kwargs.get("ids")
-        urls = kwargs.get("urls")
+        ids = kwargs.get("ids") or []
+        urls = kwargs.get("urls") or []
         input_file_path = kwargs.get("input_file_path")
         output_dir = kwargs.get("output_dir")
         language_code = kwargs.get("language_code")
         api_key = kwargs.get("api_key")
 
+        url_column_name = kwargs.get("url_column_name", cls.URL_COLUMN_NAME)
+        id_column_name = kwargs.get("id_column_name", cls.ID_COLUMN_NAME)
+
         youtube = YouTube([api_key], disable_ipv6=True)
 
-        if input_file_path:
-            ids += cls.data_from_csv(Path(input_file_path), "video_id")
+        if (input_file_path := kwargs.get("input_file_path")):
+            if (urls_from_csv := cls.data_from_csv(input_file_path, url_column_name, False)):
+                ids += [cls.video_id_from_url(url) for url in urls_from_csv]
+            if (ids_from_csv := cls.data_from_csv(input_file_path, id_column_name, False)):
+                ids += ids_from_csv
+
+        if not ids and not urls:
+            raise Exception(
+                "Either 'ids' or 'urls' must be provided for the video-transcription command"
+            )
 
         if urls:
             ids += [cls.video_id_from_url(url) for url in urls]
 
         # Remove duplicated
         ids = list(set(ids))
-        
-        # youtube.videos_transcriptions(ids, language_code, output_dir)
-
-        results = []
-        for video_id in ids:
-            try:
-                transcription = youtube.video_transcription(video_id, language_code)
-                output_file_path = cls.save_transcription_to_file(video_id, transcription, output_dir)
-                results.append(f"Transcription saved to {output_file_path}")
-            except Exception as e:
-                results.append(f"Error processing video {video_id}: {str(e)}")
-
-        return "\n".join(results)
+        youtube.videos_transcriptions(ids, language_code, output_dir)
+        output_dir_path = Path(output_dir)
+        saved_transcriptions = [
+            str(
+                output_dir_path / f"{v_id}.{language_code}.vtt"
+            ) for v_id in ids if (output_dir_path / f"{v_id}.{language_code}.vtt").is_file()
+        ]
+        return "\n".join(saved_transcriptions)