Add test_ttl

Author: caseyclements

libs/langgraph-checkpoint-mongodb/langgraph/checkpoint/mongodb/saver.py

@@ -40,7 +40,7 @@ class MongoDBSaver(BaseCheckpointSaver):
         db_name (Optional[str]): Database name
         checkpoint_collection_name (Optional[str]): Name of Collection of Checkpoints
         writes_collection_name (Optional[str]): Name of Collection of intermediate writes.
-        ttl (Optional[int]): Time to live in seconds.
+        ttl (Optional[int]): Time to live in seconds. See https://www.mongodb.com/docs/manual/core/index-ttl/.
 
     Examples:
 
@@ -87,12 +87,12 @@ def __init__(
                 keys=[("thread_id", 1), ("checkpoint_ns", 1), ("checkpoint_id", -1)],
                 unique=True,
             )
-
             if self.ttl:
                 self.checkpoint_collection.create_index(
                     keys=[("created_at", ASCENDING)],
                     expireAfterSeconds=self.ttl,
                 )
+
         if len(self.writes_collection.list_indexes().to_list()) < 2:
             self.writes_collection.create_index(
                 keys=[

libs/langgraph-checkpoint-mongodb/tests/unit_tests/test_sync.py

@@ -1,4 +1,5 @@
 import os
+from time import sleep
 from typing import Any
 
 import pytest
@@ -165,3 +166,35 @@ def test_nested_filter() -> None:
         # Drop collections
         saver.checkpoint_collection.drop()
         saver.writes_collection.drop()
+
+
+def test_ttl(input_data: dict[str, Any]) -> None:
+    collection_name = "ttl_test"
+    ttl = 1
+    background_period = 60  # period between background task runs
+
+    with MongoDBSaver.from_conn_string(
+        MONGODB_URI, DB_NAME, collection_name, ttl=ttl
+    ) as saver:
+        try:
+            # save a checkpoint
+            saver.put(
+                input_data["config_2"],
+                input_data["chkpnt_2"],
+                input_data["metadata_2"],
+                {},
+            )
+
+            query: dict[str, Any] = {}  # search by no keys, return all checkpoints
+            search_results_2 = list(saver.list(None, filter=query))
+            assert len(search_results_2) == 1
+            assert search_results_2[0].metadata == input_data["metadata_2"]
+
+            sleep(ttl + background_period)
+            assert len(list(saver.list(None, filter=query))) == 0
+
+        finally:
+            saver.checkpoint_collection.delete_many({})
+            saver.checkpoint_collection.drop_indexes()
+            saver.writes_collection.delete_many({})
+            saver.writes_collection.drop_indexes()