From 1fc19acb6cb524d836812e1776427f07d06188e4 Mon Sep 17 00:00:00 2001
From: Toni Spets <toni.spets@beeper.com>
Date: Fri, 20 Dec 2024 12:53:22 +0200
Subject: [PATCH] Test checkpoints on close will not fail

Checkpoint on close must also not force a new generation if we didn't
actually lose our lock while another process wrote to the WAL.
---
 db_test.go | 44 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/db_test.go b/db_test.go
index 4e1c4946..b9f113ce 100644
--- a/db_test.go
+++ b/db_test.go
@@ -591,6 +591,50 @@ func TestDB_Sync(t *testing.T) {
 	})
 }
 
+// Ensure a checkpoint on close will succeed.
+func TestDB_CheckpointOnClose(t *testing.T) {
+	db, sqldb := MustOpenDBs(t)
+	defer db.Close(context.Background())
+	defer MustCloseSQLDB(t, sqldb)
+
+	// Force checkpoint after every write.
+	db.MinCheckpointPageN = 1
+
+	// Execute a query to force a write to the WAL and then sync.
+	if _, err := sqldb.Exec(`CREATE TABLE foo (bar TEXT);`); err != nil {
+		t.Fatal(err)
+	} else if err := db.Sync(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	// Get generation to test later it didn't change.
+	firstGeneration, err := db.CurrentGeneration()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Write again and sync.
+	if _, err := sqldb.Exec(`INSERT INTO foo (bar) VALUES ('baz');`); err != nil {
+		t.Fatal(err)
+	} else if err := db.Sync(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	// Write to WAL & close.
+	if _, err := sqldb.Exec(`INSERT INTO foo (bar) VALUES ('baz');`); err != nil {
+		t.Fatal(err)
+	} else if err := db.Close(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	finalGeneration, err := db.CurrentGeneration()
+	if err != nil {
+		t.Fatal(err)
+	} else if finalGeneration != firstGeneration {
+		t.Fatal("generation changed", firstGeneration, "!=", finalGeneration)
+	}
+}
+
 // MustOpenDBs returns a new instance of a DB & associated SQL DB.
 func MustOpenDBs(tb testing.TB) (*litestream.DB, *sql.DB) {
 	tb.Helper()