diff --git a/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py b/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py
index 7c79ad01..d2373dc0 100644
--- a/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py
+++ b/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py
@@ -214,6 +214,7 @@ def test_vector_search_retriever_tool_description_generation(index_name: str) ->
 def test_vs_tool_tracing(index_name: str, tool_name: str | None) -> None:
     vector_search_tool = init_vector_search_tool(index_name, tool_name=tool_name)
     vector_search_tool._run("Databricks Agent Framework")
+    mlflow.flush_trace_async_logging()
     trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
     spans = trace.search_spans(name=tool_name or index_name, span_type=SpanType.RETRIEVER)
     assert len(spans) == 1
diff --git a/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py b/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py
index a87698c6..459130f0 100644
--- a/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py
+++ b/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py
@@ -187,6 +187,7 @@ def test_vector_search_retriever_tool_init(
     assert all(["id" in d["metadata"] for d in docs])
 
     # Ensure tracing works properly
+    mlflow.flush_trace_async_logging()
     trace = mlflow.get_trace(mlflow.get_last_active_trace_id())
     spans = trace.search_spans(name=tool_name or index_name, span_type=SpanType.RETRIEVER)
     assert len(spans) == 1