diff --git a/tests/test_auto_llm.py b/tests/test_auto_llm.py
index c12278e1..5a30ec40 100644
--- a/tests/test_auto_llm.py
+++ b/tests/test_auto_llm.py
@@ -1,5 +1,6 @@
 from llama_index import Document, ServiceContext, VectorStoreIndex
 from llama_index.llms.base import LLM
+from llama_index.query_engine import BaseQueryEngine
 
 from autollm.auto.llm import AutoLLM
 
@@ -18,7 +19,5 @@ def test_auto_llm():
 
     query_engine = index.as_query_engine()
 
-    response = query_engine.query("What is the meaning of life?")
-
-    # Check if the response is not None
-    assert response.response is not None
+    # Check if the query_engine is an instance of BaseQueryEngine
+    assert isinstance(query_engine, BaseQueryEngine)
diff --git a/tests/test_auto_query_engine.py b/tests/test_auto_query_engine.py
index a40488d2..6e87fe4e 100644
--- a/tests/test_auto_query_engine.py
+++ b/tests/test_auto_query_engine.py
@@ -1,14 +1,13 @@
 from llama_index import Document
+from llama_index.query_engine import BaseQueryEngine
 
 from autollm.auto.query_engine import AutoQueryEngine
 
 
 def test_auto_query_engine():
-    documents = Document.example()
+    documents = [Document.example()]
     vector_store_params = {"vector_store_type": "VectorStoreIndex", "documents": documents}
     query_engine = AutoQueryEngine.from_parameters(vector_store_params=vector_store_params)
 
-    response = query_engine.query("What is the meaning of life?")
-
-    # Check if the response is not None
-    assert response.response is not None
+    # Check if the query_engine is an instance of BaseQueryEngine
+    assert isinstance(query_engine, BaseQueryEngine)
diff --git a/tests/test_auto_service_context.py b/tests/test_auto_service_context.py
index be57dd47..17b382a4 100644
--- a/tests/test_auto_service_context.py
+++ b/tests/test_auto_service_context.py
@@ -1,4 +1,4 @@
-from llama_index import Document, VectorStoreIndex
+from llama_index import Document, ServiceContext, VectorStoreIndex
 
 from autollm.auto.service_context import AutoServiceContext
 
@@ -8,6 +8,9 @@ def test_auto_service_context():
 
     service_context = AutoServiceContext.from_defaults(enable_cost_calculator=True)
 
+    # Check if the service_context is an instance of ServiceContext
+    assert isinstance(service_context, ServiceContext)
+
     index = VectorStoreIndex.from_documents(documents=[document], service_context=service_context)
 
     query_engine = index.as_query_engine()
@@ -17,6 +20,6 @@ def test_auto_service_context():
     # Check if the response is not None
     assert response.response is not None
 
-    # Check if the total token cost is greater than 0
+    # Check if the cost calculating handler is working
     cost_caltulator = service_context.callback_manager.handlers[0]
     assert cost_caltulator.total_llm_token_cost > 0
diff --git a/tests/test_auto_vector_store_index.py b/tests/test_auto_vector_store_index.py
index 17198897..963ee480 100644
--- a/tests/test_auto_vector_store_index.py
+++ b/tests/test_auto_vector_store_index.py
@@ -1,10 +1,11 @@
 from llama_index import Document, VectorStoreIndex
+from llama_index.query_engine import BaseQueryEngine
 
 from autollm.auto.vector_store_index import AutoVectorStoreIndex
 
 
 def test_auto_vector_store():
-    documents = Document.example()
+    documents = [Document.example()]
 
     vector_store_index = AutoVectorStoreIndex.from_defaults(
         vector_store_type="VectorStoreIndex", documents=documents)
@@ -14,7 +15,5 @@ def test_auto_vector_store():
 
     query_engine = vector_store_index.as_query_engine()
 
-    response = query_engine.query("What is the meaning of life?")
-
-    # Check if the response is not None
-    assert response.response is not None
+    # Check if the query_engine is an instance of BaseQueryEngine
+    assert isinstance(query_engine, BaseQueryEngine)
diff --git a/tests/test_cost_calculating_handler.py b/tests/test_cost_calculating_handler.py
deleted file mode 100644
index 1a0d8774..00000000
--- a/tests/test_cost_calculating_handler.py
+++ /dev/null
@@ -1,23 +0,0 @@
-from llama_index import Document, ServiceContext, VectorStoreIndex
-from llama_index.callbacks import CallbackManager
-
-from autollm.callbacks.cost_calculating import CostCalculatingHandler
-
-
-def test_cost_calculating_handler():
-    document = Document.example()
-
-    cost_calculater = CostCalculatingHandler(model="gpt-3.5-turbo", verbose=False)
-    callback_manager = CallbackManager([cost_calculater])
-    service_context = ServiceContext.from_defaults(callback_manager=callback_manager)
-
-    index = VectorStoreIndex.from_documents(documents=[document], service_context=service_context)
-
-    query_engine = index.as_query_engine()
-
-    response = query_engine.query("What is the meaning of life?")
-    # Check if the response is not None
-    assert response.response is not None
-
-    # Check if the total token cost is greater than 0
-    assert cost_calculater.total_llm_token_cost > 0