Merge pull request #354 from tejhande/patch-2

Enhance tests for FetchNode with mocking

Author: VinciGit00

tests/nodes/fetch_node_test.py

@@ -1,107 +1,104 @@
 import os
 import pytest
+from unittest.mock import patch, MagicMock
 from scrapegraphai.nodes import FetchNode
 
-def test_fetch_node_html():
+def get_file_path(file_name):
     """
-    Run the tests
+    Helper function to get the absolute file path.
     """
+    curr_dir = os.path.dirname(os.path.realpath(__file__))
+    file_path = os.path.join(curr_dir, file_name)
+    return file_path
+
+@patch('scrapegraphai.nodes.FetchNode.execute')
+def test_fetch_node_html(mock_execute):
+    """
+    Test FetchNode with HTML input.
+    """
+    mock_execute.return_value = MagicMock()
     fetch_node = FetchNode(
         input="url | local_dir",
         output=["doc"],
         node_config={
             "headless": False
         }
     )
-
     state = {
         "url": "https://twitter.com/home"
     }
-
     result = fetch_node.execute(state)
-
     assert result is not None
+    mock_execute.assert_called_once_with(state)
 
-def test_fetch_node_json():
+@patch('scrapegraphai.nodes.FetchNode.execute')
+def test_fetch_node_json(mock_execute):
     """
-    Run the tests
+    Test FetchNode with JSON input.
     """
-    FILE_NAME_JSON = "inputs/example.json"
-    curr_dir = os.path.dirname(os.path.realpath(__file__))
-    file_path_json = os.path.join(curr_dir, FILE_NAME_JSON)
-
+    mock_execute.return_value = MagicMock()
+    file_path_json = get_file_path("inputs/example.json")
     state_json = {
         "json": file_path_json
     }
-
     fetch_node_json = FetchNode(
         input="json",
         output=["doc"],
     )
-
     result_json = fetch_node_json.execute(state_json)
-
     assert result_json is not None
+    mock_execute.assert_called_once_with(state_json)
 
-def test_fetch_node_xml():
+@patch('scrapegraphai.nodes.FetchNode.execute')
+def test_fetch_node_xml(mock_execute):
     """
-    Run the tests
+    Test FetchNode with XML input.
     """
-    FILE_NAME_XML = "inputs/books.xml"
-    curr_dir = os.path.dirname(os.path.realpath(__file__))
-    file_path_xml = os.path.join(curr_dir, FILE_NAME_XML)
-
+    mock_execute.return_value = MagicMock()
+    file_path_xml = get_file_path("inputs/books.xml")
     state_xml = {
         "xml": file_path_xml
     }
-
     fetch_node_xml = FetchNode(
         input="xml",
         output=["doc"],
     )
-
     result_xml = fetch_node_xml.execute(state_xml)
-
     assert result_xml is not None
+    mock_execute.assert_called_once_with(state_xml)
 
-def test_fetch_node_csv():
+@patch('scrapegraphai.nodes.FetchNode.execute')
+def test_fetch_node_csv(mock_execute):
     """
-    Run the tests
+    Test FetchNode with CSV input.
     """
-    FILE_NAME_CSV = "inputs/username.csv"
-    curr_dir = os.path.dirname(os.path.realpath(__file__))
-    file_path_csv = os.path.join(curr_dir, FILE_NAME_CSV)
-
+    mock_execute.return_value = MagicMock()
+    file_path_csv = get_file_path("inputs/username.csv")
     state_csv = {
-        "csv": file_path_csv  # Definire un dizionario con la chiave "csv" e il valore come percorso del file CSV
+        "csv": file_path_csv
     }
-
     fetch_node_csv = FetchNode(
         input="csv",
         output=["doc"],
     )
-
     result_csv = fetch_node_csv.execute(state_csv)
-
     assert result_csv is not None
+    mock_execute.assert_called_once_with(state_csv)
 
-def test_fetch_node_txt():
+@patch('scrapegraphai.nodes.FetchNode.execute')
+def test_fetch_node_txt(mock_execute):
     """
-    Run the tests
+    Test FetchNode with TXT input.
     """
-    FILE_NAME_TXT = "inputs/plain_html_example.txt"
-    curr_dir = os.path.dirname(os.path.realpath(__file__))
-    file_path_txt = os.path.join(curr_dir, FILE_NAME_TXT)
-
+    mock_execute.return_value = MagicMock()
+    file_path_txt = get_file_path("inputs/plain_html_example.txt")
     state_txt = {
-        "txt": file_path_txt  # Definire un dizionario con la chiave "txt" e il valore come percorso del file TXT
+        "txt": file_path_txt
     }
-
     fetch_node_txt = FetchNode(
         input="txt",
         output=["doc"],
     )
-
     result_txt = fetch_node_txt.execute(state_txt)
-
     assert result_txt is not None
+    mock_execute.assert_called_once_with(state_txt)