diff --git a/dsLightRag/Test/T1_WriteToPg.py b/dsLightRag/Test/T1_WriteToPg.py
index 0a33cf2d..b8c32e29 100644
--- a/dsLightRag/Test/T1_WriteToPg.py
+++ b/dsLightRag/Test/T1_WriteToPg.py
@@ -63,9 +63,13 @@ async def initialize_pg_rag():
 
 
 async def main():
-    rag = await initialize_pg_rag()
-    with open(f"../Txt/sushi.txt", "r", encoding="utf-8") as f:
-        await rag.ainsert(f.read())
+    try:
+        rag = await initialize_pg_rag()
+        with open(f"../Txt/sushi.txt", "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+    finally:
+        if rag:
+            await rag.finalize_storages()
 
 if __name__ == "__main__":
     asyncio.run(main())
diff --git a/dsLightRag/Test/T2_ReadFromPg.py b/dsLightRag/Test/T2_ReadFromPg.py
index 377d0b14..0ce130d6 100644
--- a/dsLightRag/Test/T2_ReadFromPg.py
+++ b/dsLightRag/Test/T2_ReadFromPg.py
@@ -3,12 +3,12 @@ import inspect
 import logging
 import os
 
+from lightrag import QueryParam, LightRAG
 from lightrag.kg.shared_storage import initialize_pipeline_status
 from lightrag.utils import EmbeddingFunc
 
 from Config.Config import LLM_MODEL_NAME, EMBED_DIM, EMBED_MAX_TOKEN_SIZE
-from Util.LightRagUtil import configure_logging, initialize_rag, print_stream, llm_model_func, embedding_func
-from lightrag import QueryParam, LightRAG
+from Util.LightRagUtil import configure_logging, print_stream, llm_model_func, embedding_func
 
 # 在程序开始时添加以下配置
 logging.basicConfig(