diff --git a/tests/test_loading.py b/tests/test_loading.py
index a1a2ec1f4..74aba58db 100644
--- a/tests/test_loading.py
+++ b/tests/test_loading.py
@@ -1,17 +1,14 @@
-import json
-
 import pytest
-from langchain.chains.base import Chain
+
 from langflow.graph import Graph
 from langflow.processing.load import load_flow_from_json
-from langflow.utils.payload import get_root_vertex
 
 
 def test_load_flow_from_json():
     """Test loading a flow from a json file"""
     loaded = load_flow_from_json(pytest.BASIC_EXAMPLE_PATH)
     assert loaded is not None
-    assert isinstance(loaded, Chain)
+    assert isinstance(loaded, Graph)
 
 
 def test_load_flow_from_json_with_tweaks():
@@ -19,18 +16,5 @@ def test_load_flow_from_json_with_tweaks():
     tweaks = {"dndnode_82": {"model_name": "gpt-3.5-turbo-16k-0613"}}
     loaded = load_flow_from_json(pytest.BASIC_EXAMPLE_PATH, tweaks=tweaks)
     assert loaded is not None
-    assert isinstance(loaded, Chain)
+    assert isinstance(loaded, Graph)
     assert loaded.llm.model_name == "gpt-3.5-turbo-16k-0613"
-
-
-def test_get_root_vertex():
-    with open(pytest.BASIC_EXAMPLE_PATH, "r") as f:
-        flow_graph = json.load(f)
-    data_graph = flow_graph["data"]
-    nodes = data_graph["nodes"]
-    edges = data_graph["edges"]
-    graph = Graph(nodes, edges)
-    root = get_root_vertex(graph)
-    assert root is not None
-    assert hasattr(root, "id")
-    assert hasattr(root, "data")