From 9cd516507cc5ad65b100522b488cb0272dc7b366 Mon Sep 17 00:00:00 2001
From: VinciGit00 <mvincig11@gmail.com>
Date: Sat, 27 Apr 2024 21:23:35 +0200
Subject: [PATCH] fix: bug with fetch node

---
 examples/single_node/fetch_node.py | 27 ++++++++++++++++++
 scrapegraphai/nodes/fetch_node.py  |  4 +--
 tests/nodes/.env.example           |  1 -
 tests/nodes/fetch_node_test.py     | 44 ++++++++++++++++++++++++++++++
 tests/nodes/robot_node_test.py     | 14 ++--------
 5 files changed, 76 insertions(+), 14 deletions(-)
 delete mode 100644 tests/nodes/.env.example
 create mode 100644 tests/nodes/fetch_node_test.py

diff --git a/examples/single_node/fetch_node.py b/examples/single_node/fetch_node.py
index e69de29b..90660996 100644
--- a/examples/single_node/fetch_node.py
+++ b/examples/single_node/fetch_node.py
@@ -0,0 +1,27 @@
+"""
+Example of custom graph using existing nodes
+"""
+
+from scrapegraphai.nodes import FetchNode
+
+# ************************************************
+# Define the node
+# ************************************************
+
+
+robots_node = FetchNode(
+    input="url | local_dir",
+    output=["doc"],
+)
+
+# ************************************************
+# Test the node
+# ************************************************
+
+state = {
+    "url": "https://twitter.com/home"
+}
+
+result = robots_node.execute(state)
+
+print(result)
diff --git a/scrapegraphai/nodes/fetch_node.py b/scrapegraphai/nodes/fetch_node.py
index 51fc7c30..76d80929 100644
--- a/scrapegraphai/nodes/fetch_node.py
+++ b/scrapegraphai/nodes/fetch_node.py
@@ -59,7 +59,7 @@ class FetchNode(BaseNode):
 
         Raises:
             KeyError: If the 'url' key is not found in the state, indicating that the
-                      necessary information to perform the operation is missing.
+                    necessary information to perform the operation is missing.
         """
         print(f"--- Executing {self.node_name} Node ---")
 
@@ -78,7 +78,7 @@ class FetchNode(BaseNode):
             })]
 
         else:
-            if self.node_config.get("endpoint") is not None:
+            if self.node_config is not None and self.node_config.get("endpoint") is not None:
                 loader = AsyncHtmlLoader(
                     source, proxies={"http": self.node_config["endpoint"]})
             else:
diff --git a/tests/nodes/.env.example b/tests/nodes/.env.example
deleted file mode 100644
index 12c1491c..00000000
--- a/tests/nodes/.env.example
+++ /dev/null
@@ -1 +0,0 @@
-OPENAI_APIKEY="your openai api key"
\ No newline at end of file
diff --git a/tests/nodes/fetch_node_test.py b/tests/nodes/fetch_node_test.py
new file mode 100644
index 00000000..e0552a05
--- /dev/null
+++ b/tests/nodes/fetch_node_test.py
@@ -0,0 +1,44 @@
+"""
+Module for testinh robot_node
+"""
+import pytest
+from scrapegraphai.nodes import FetchNode
+
+
+@pytest.fixture
+def setup():
+    """
+    setup
+    """
+    # ************************************************
+    # Define the node
+    # ************************************************
+
+    robots_node = FetchNode(
+        input="url | local_dir",
+        output=["doc"],
+    )
+
+    return robots_node
+
+# ************************************************
+# Test the node
+# ************************************************
+
+
+def test_robots_node(setup):
+    """
+    Run the tests
+    """
+    state = {
+        "url": "https://twitter.com/home"
+    }
+
+    result = setup.execute(state)
+
+    assert result is not None
+
+
+# If you need to run this script directly
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/nodes/robot_node_test.py b/tests/nodes/robot_node_test.py
index 3d57e60c..7808a976 100644
--- a/tests/nodes/robot_node_test.py
+++ b/tests/nodes/robot_node_test.py
@@ -1,15 +1,10 @@
 """
 Module for testinh robot_node
 """
-import os
-from dotenv import load_dotenv
 import pytest
-from scrapegraphai.models import OpenAI
+from scrapegraphai.models import Ollama
 from scrapegraphai.nodes import RobotsNode
 
-# Load environment variables from .env file
-load_dotenv()
-
 
 @pytest.fixture
 def setup():
@@ -20,12 +15,9 @@ def setup():
     # Define the configuration for the graph
     # ************************************************
 
-    openai_key = os.getenv("OPENAI_APIKEY")
-
     graph_config = {
         "llm": {
-            "api_key": openai_key,
-            "model": "gpt-3.5-turbo",
+            "model": "ollama/llama3",
             "temperature": 0,
             "streaming": True
         },
@@ -35,7 +27,7 @@ def setup():
     # Define the node
     # ************************************************
 
-    llm_model = OpenAI(graph_config["llm"])
+    llm_model = Ollama(graph_config["llm"])
 
     robots_node = RobotsNode(
         input="url",