def test_custom_ttl_attribute(chat_store_with_custom_ttl):
    """Test that custom TTL attribute name is used."""
    messages = [ChatMessage(content="Custom TTL Attribute")]
    chat_store_with_custom_ttl.set_messages("CustomTTLSession", messages)

    # Get the raw item from DynamoDB to check TTL
    response = chat_store_with_custom_ttl._table.get_item(
        Key={"SessionId": "CustomTTLSession"}
    )
    item = response["Item"]

    # Verify custom TTL attribute exists
    assert "ExpiresAt" in item
    assert "TTL" not in item

    # Verify TTL is set to a future time (now + 7200 seconds)
    current_time = int(time.time())
    assert item["ExpiresAt"] > current_time
    # Allow some buffer for test execution time
    assert item["ExpiresAt"] <= current_time + 7300
Test that custom TTL attribute name is used.
llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/tests/test_chat_store_dynamodb_chat_store.py