From ffb2aad45d55d120db04709b9673c1aa955c077c Mon Sep 17 00:00:00 2001
From: Gabriel Almeida <gabriel@logspace.ai>
Date: Sat, 29 Apr 2023 07:38:06 -0300
Subject: [PATCH] fixes to tests

---
 src/backend/langflow/template/nodes.py |  1 +
 tests/test_agents_template.py          |  2 +-
 tests/test_chains_template.py          | 14 ++++++++++----
 tests/test_graph.py                    |  1 -
 tests/test_websocket.py                |  3 ---
 5 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/backend/langflow/template/nodes.py b/src/backend/langflow/template/nodes.py
index acf93c948..7b5ac6d3f 100644
--- a/src/backend/langflow/template/nodes.py
+++ b/src/backend/langflow/template/nodes.py
@@ -199,6 +199,7 @@ class SeriesCharacterChainNode(FrontendNode):
         "Chain",
         "ConversationChain",
         "SeriesCharacterChain",
+        "function",
     ]
 
 
diff --git a/tests/test_agents_template.py b/tests/test_agents_template.py
index a546454eb..4ed9d1153 100644
--- a/tests/test_agents_template.py
+++ b/tests/test_agents_template.py
@@ -130,7 +130,7 @@ def test_initialize_agent(client: TestClient):
     agents = json_response["agents"]
 
     initialize_agent = agents["initialize_agent"]
-    assert initialize_agent["base_classes"] == ["AgentExecutor"]
+    assert initialize_agent["base_classes"] == ["AgentExecutor", "function"]
     template = initialize_agent["template"]
 
     assert template["agent"] == {
diff --git a/tests/test_chains_template.py b/tests/test_chains_template.py
index cff844c90..13bef2e9c 100644
--- a/tests/test_chains_template.py
+++ b/tests/test_chains_template.py
@@ -20,7 +20,12 @@ def test_conversation_chain(client: TestClient):
     chain = chains["ConversationChain"]
 
     # Test the base classes, template, memory, verbose, llm, input_key, output_key, and _type objects
-    assert set(chain["base_classes"]) == {"LLMChain", "ConversationChain", "Chain"}
+    assert set(chain["base_classes"]) == {
+        "function",
+        "LLMChain",
+        "ConversationChain",
+        "Chain",
+    }
     template = chain["template"]
     assert template["memory"] == {
         "required": False,
@@ -96,7 +101,7 @@ def test_llm_chain(client: TestClient):
     chain = chains["LLMChain"]
 
     # Test the base classes, template, memory, verbose, llm, input_key, output_key, and _type objects
-    assert set(chain["base_classes"]) == {"LLMChain", "Chain"}
+    assert set(chain["base_classes"]) == {"function", "LLMChain", "Chain"}
     template = chain["template"]
     assert template["memory"] == {
         "required": False,
@@ -154,7 +159,7 @@ def test_llm_checker_chain(client: TestClient):
     chain = chains["LLMCheckerChain"]
 
     # Test the base classes, template, memory, verbose, llm, input_key, output_key, and _type objects
-    assert set(chain["base_classes"]) == {"LLMCheckerChain", "Chain"}
+    assert set(chain["base_classes"]) == {"function", "LLMCheckerChain", "Chain"}
     template = chain["template"]
     assert template["memory"] == {
         "required": False,
@@ -231,7 +236,7 @@ def test_llm_math_chain(client: TestClient):
     chain = chains["LLMMathChain"]
 
     # Test the base classes, template, memory, verbose, llm, input_key, output_key, and _type objects
-    assert set(chain["base_classes"]) == {"LLMMathChain", "Chain"}
+    assert set(chain["base_classes"]) == {"function", "LLMMathChain", "Chain"}
     template = chain["template"]
     assert template["memory"] == {
         "required": False,
@@ -310,6 +315,7 @@ def test_series_character_chain(client: TestClient):
 
     # Test the base classes, template, memory, verbose, llm, input_key, output_key, and _type objects
     assert set(chain["base_classes"]) == {
+        "function",
         "LLMChain",
         "BaseCustomChain",
         "Chain",
diff --git a/tests/test_graph.py b/tests/test_graph.py
index 76451fe6a..e109850e3 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -1,4 +1,3 @@
-import json
 from typing import Type, Union
 
 import pytest
diff --git a/tests/test_websocket.py b/tests/test_websocket.py
index 1137c4813..7ac646cf4 100644
--- a/tests/test_websocket.py
+++ b/tests/test_websocket.py
@@ -1,6 +1,5 @@
 import json
 from unittest.mock import patch
-from langflow.api.schemas import ChatMessage
 from fastapi.testclient import TestClient
 
 
@@ -11,8 +10,6 @@ def test_websocket_connection(client: TestClient):
 
 
 def test_chat_history(client: TestClient):
-    chat_history = []
-
     # Mock the process_graph function to return a specific value
     with patch("langflow.api.chat_manager.process_graph") as mock_process_graph:
         mock_process_graph.return_value = ("Hello, I'm a mock response!", "")