- Add test for video_search command;

- Add updates to some test files;
- created conftest file

Author: aninhasalesp

tests/commands/conftest.py

@@ -0,0 +1,29 @@
+import pytest
+
+
+@pytest.fixture
+def channels_urls():
+    return [
+        "https://www.youtube.com/@Turicas/featured",
+        "https://www.youtube.com/c/PythonicCaf%C3%A9"
+    ]
+
+
+@pytest.fixture
+def videos_ids():
+    return [
+        "video_id_1",
+        "video_id_2"
+    ]
+
+
+@pytest.fixture
+def videos_urls(videos_ids):
+    return [
+        f"https://www.youtube.com/?v={video_id}" for video_id in videos_ids
+    ]
+
+
+@pytest.fixture
+def usernames():
+    return ["Turicas", "PythonicCafe"]

tests/commands/test_channel_info.py

@@ -52,10 +52,10 @@ def test_channel_ids_from_urls_and_usernames(mocker):
     youtube_mock.return_value.channel_id_from_username = channel_id_from_username_mock
     youtube_mock.return_value.channels_infos = channels_infos_mock
 
-    ChannelInfo.execute(urls=urls, usernames=usernames)
+    ChannelInfo.execute(urls=channels_urls, usernames=usernames)
 
     channel_id_from_url_mock.assert_has_calls(
-        [call(url) for url in urls]
+        [call(url) for url in channels_urls]
     )
     channel_id_from_username_mock.assert_has_calls(
         [call(username) for username in usernames]

tests/commands/test_video_search.py

@@ -0,0 +1,66 @@
+import csv
+import pytest
+
+from io import StringIO
+from unittest.mock import Mock, call
+from datetime import datetime
+
+from youtool.commands.video_search import VideoSearch
+
+
+def test_video_search_string_output(mocker, videos_ids, videos_urls):
+    youtube_mock = mocker.patch("youtool.commands.video_search.YouTube")
+    expected_videos_infos = [
+        {
+            column: f"v_{index}" for column in VideoSearch.INFO_COLUMNS
+        } for index, _ in enumerate(videos_ids)
+    ]
+
+    csv_file = StringIO()
+    csv_writer = csv.DictWriter(csv_file, fieldnames=VideoSearch.INFO_COLUMNS)
+    csv_writer.writeheader()
+    csv_writer.writerows(expected_videos_infos)
+
+    videos_infos_mock = Mock(return_value=expected_videos_infos)
+    youtube_mock.return_value.videos_infos = videos_infos_mock
+
+    result = VideoSearch.execute(ids=videos_ids, urls=videos_urls)
+
+    videos_infos_mock.assert_called_once_with(list(set(videos_ids)))
+    assert result == csv_file.getvalue()
+
+
+def test_video_search_file_output(mocker, videos_ids, videos_urls, tmp_path):
+    youtube_mock = mocker.patch("youtool.commands.video_search.YouTube")
+    expected_videos_infos = [
+        {
+            column: f"v_{index}" for column in VideoSearch.INFO_COLUMNS
+        } for index, _ in enumerate(videos_ids)
+    ]
+
+    expected_csv_file = StringIO()
+    csv_writer = csv.DictWriter(expected_csv_file, fieldnames=VideoSearch.INFO_COLUMNS)
+    csv_writer.writeheader()
+    csv_writer.writerows(expected_videos_infos)
+
+    timestamp = datetime.now().strftime("%f")
+    output_file_name = f"output_{timestamp}.csv"
+    output_file_path = tmp_path / output_file_name
+
+    videos_infos_mock = Mock(return_value=expected_videos_infos)
+    youtube_mock.return_value.videos_infos = videos_infos_mock
+
+    result_file_path = VideoSearch.execute(
+        ids=videos_ids, urls=videos_urls, output_file_path=output_file_path
+    )
+
+    with open(result_file_path, "r") as result_csv_file:
+        result_csv = result_csv_file.read()
+
+    videos_infos_mock.assert_called_once_with(list(set(videos_ids)))
+    assert result_csv.replace("\r", "") == expected_csv_file.getvalue().replace("\r", "")
+
+
+def test_video_search_no_id_and_url_error():
+    with pytest.raises(Exception, match="Either 'ids' or 'urls' must be provided"):
+        VideoSearch.execute(ids=None, urls=None)

youtool/commands/base.py

@@ -105,7 +105,7 @@ def data_from_csv(file_path: Path, data_column_name: Optional[str] = None) -> Li
 
             if fieldnames is None:
                 raise ValueError("Fieldnames is None")
-            
+
             if data_column_name not in fieldnames:
                 raise Exception(f"Column {data_column_name} not found on {file_path}")
             for row in reader:

youtool/commands/video_search.py

@@ -8,15 +8,19 @@
 
 
 class VideoSearch(Command):
-    """Search video info from a list of IDs or URLs (or CSV filename with URLs/IDs inside), generate CSV output (simplified video dict schema or option to get full video info)
+    """
+    Search video info from a list of IDs or URLs (or CSV filename with URLs/IDs inside),
+    generate CSV output (simplified video dict schema or option to get full video info)
     """
     name = "video-search"
     arguments = [
         {"name": "--ids", "type": str, "help": "Video IDs", "nargs": "*"},
         {"name": "--urls", "type": str, "help": "Video URLs", "nargs": "*"},
         {"name": "--input-file-path", "type": str, "help": "Input CSV file path with URLs/IDs"},
         {"name": "--output-file-path", "type": str, "help": "Output CSV file path"},
-        {"name": "--full-info", "type": bool, "help": "Option to get full video info", "default": False}
+        {"name": "--full-info", "type": bool, "help": "Option to get full video info", "default": False},
+        {"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"}
     ]
 
     ID_COLUMN_NAME: str = "video_id"
@@ -28,23 +32,6 @@ class VideoSearch(Command):
         "id", "title", "description", "published_at", "view_count", "like_count", "comment_count"
     ]
 
-    @staticmethod
-    def filter_fields(video_info: Dict, info_columns: Optional[List] = None) -> Dict:
-        """Filters the fields of a dictionary containing video information based on specified columns.
-
-        Args:
-            video_info (Dict): A dictionary containing video information.
-            info_columns (Optional[List], optional): A list specifying which fields to include in the filtered output. 
-            If None, returns the entire video_info dictionary. Defaults to None.
-
-        Returns:
-            A dictionary containing only the fields specified in info_columns (if provided) 
-            or the entire video_info dictionary if info_columns is None.
-        """
-        return {
-            field: value for field, value in video_info.items() if field in info_columns
-        } if info_columns else video_info
-
     @classmethod
     def execute(cls: Self, **kwargs) -> str:
         """
@@ -68,21 +55,20 @@ def execute(cls: Self, **kwargs) -> str:
         """
         ids = kwargs.get("ids", [])
         urls = kwargs.get("urls", [])
-        input_file_path = kwargs.get("input_file_path")
         output_file_path = kwargs.get("output_file_path")
         api_key = kwargs.get("api_key")
         full_info = kwargs.get("full_info", False)
 
+        url_column_name = kwargs.get("url_column_name", cls.URL_COLUMN_NAME)
+        id_column_name = kwargs.get("id_column_name", cls.ID_COLUMN_NAME)
+
         info_columns = VideoSearch.FULL_INFO_COLUMNS if full_info else VideoSearch.INFO_COLUMNS
 
-        if input_file_path:
-            with open(input_file_path, mode='r') as infile:
-                reader = csv.DictReader(infile)
-                for row in reader:
-                    if cls.ID_COLUMN_NAME in row and row[cls.ID_COLUMN_NAME]:
-                        ids.append(row[cls.ID_COLUMN_NAME])
-                    elif cls.URL_COLUMN_NAME in row and row[cls.URL_COLUMN_NAME]:
-                        urls.append(row[cls.URL_COLUMN_NAME])
+        if (input_file_path := kwargs.get("input_file_path")):
+            if (urls_from_csv := cls.data_from_csv(input_file_path, url_column_name)):
+                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)):
+                ids += ids_from_csv
 
         if not ids and not urls:
             raise Exception("Either 'ids' or 'urls' must be provided for the video-search command")
@@ -95,7 +81,7 @@ def execute(cls: Self, **kwargs) -> str:
         # Remove duplicated
         ids = list(set(ids))
         videos_infos = list(youtube.videos_infos([_id for _id in ids if _id]))
-                                     
+
         return cls.data_to_csv(
             data=[
                 VideoSearch.filter_fields(