diff --git a/pkg/cli/context/files.go b/pkg/cli/context/files.go
new file mode 100644
index 00000000..098b2cd5
--- /dev/null
+++ b/pkg/cli/context/files.go
@@ -0,0 +1,51 @@
+/* Copyright (C) 2019, 2020, 2021, 2022, 2023, 2024, 2025 Dnote contributors
+ *
+ * This file is part of Dnote.
+ *
+ * Dnote is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * Dnote is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Dnote.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package context
+
+import (
+	"path/filepath"
+
+	"github.com/dnote/dnote/pkg/cli/consts"
+	"github.com/dnote/dnote/pkg/cli/utils"
+	"github.com/pkg/errors"
+)
+
+// InitDnoteDirs creates the dnote directories if they don't already exist.
+func InitDnoteDirs(paths Paths) error {
+	if paths.Config != "" {
+		configDir := filepath.Join(paths.Config, consts.DnoteDirName)
+		if err := utils.EnsureDir(configDir); err != nil {
+			return errors.Wrap(err, "initializing config dir")
+		}
+	}
+	if paths.Data != "" {
+		dataDir := filepath.Join(paths.Data, consts.DnoteDirName)
+		if err := utils.EnsureDir(dataDir); err != nil {
+			return errors.Wrap(err, "initializing data dir")
+		}
+	}
+	if paths.Cache != "" {
+		cacheDir := filepath.Join(paths.Cache, consts.DnoteDirName)
+		if err := utils.EnsureDir(cacheDir); err != nil {
+			return errors.Wrap(err, "initializing cache dir")
+		}
+	}
+
+	return nil
+}
diff --git a/pkg/cli/context/files_test.go b/pkg/cli/context/files_test.go
new file mode 100644
index 00000000..2422a795
--- /dev/null
+++ b/pkg/cli/context/files_test.go
@@ -0,0 +1,65 @@
+/* Copyright (C) 2019, 2020, 2021, 2022, 2023, 2024, 2025 Dnote contributors
+ *
+ * This file is part of Dnote.
+ *
+ * Dnote is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * Dnote is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Dnote.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package context
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/dnote/dnote/pkg/assert"
+	"github.com/dnote/dnote/pkg/cli/consts"
+)
+
+func assertDirsExist(t *testing.T, paths Paths) {
+	configDir := filepath.Join(paths.Config, consts.DnoteDirName)
+	info, err := os.Stat(configDir)
+	assert.Equal(t, err, nil, "config dir should exist")
+	assert.Equal(t, info.IsDir(), true, "config should be a directory")
+
+	dataDir := filepath.Join(paths.Data, consts.DnoteDirName)
+	info, err = os.Stat(dataDir)
+	assert.Equal(t, err, nil, "data dir should exist")
+	assert.Equal(t, info.IsDir(), true, "data should be a directory")
+
+	cacheDir := filepath.Join(paths.Cache, consts.DnoteDirName)
+	info, err = os.Stat(cacheDir)
+	assert.Equal(t, err, nil, "cache dir should exist")
+	assert.Equal(t, info.IsDir(), true, "cache should be a directory")
+}
+
+func TestInitDnoteDirs(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	paths := Paths{
+		Config: filepath.Join(tmpDir, "config"),
+		Data:   filepath.Join(tmpDir, "data"),
+		Cache:  filepath.Join(tmpDir, "cache"),
+	}
+
+	// Initialize directories
+	err := InitDnoteDirs(paths)
+	assert.Equal(t, err, nil, "InitDnoteDirs should succeed")
+	assertDirsExist(t, paths)
+
+	// Call again - should be idempotent
+	err = InitDnoteDirs(paths)
+	assert.Equal(t, err, nil, "InitDnoteDirs should succeed when dirs already exist")
+	assertDirsExist(t, paths)
+}
diff --git a/pkg/cli/context/testutils.go b/pkg/cli/context/testutils.go
index 17d11847..62fed833 100644
--- a/pkg/cli/context/testutils.go
+++ b/pkg/cli/context/testutils.go
@@ -19,7 +19,6 @@
 package context
 
 import (
-	"os"
 	"path/filepath"
 	"testing"
 
@@ -40,34 +39,16 @@ func getDefaultTestPaths(t *testing.T) Paths {
 	}
 }
 
-// createTestDirectories creates test directories for the given paths
-func createTestDirectories(t *testing.T, paths Paths) {
-	if paths.Config != "" {
-		configDir := filepath.Join(paths.Config, consts.DnoteDirName)
-		if err := os.MkdirAll(configDir, 0755); err != nil {
-			t.Fatal(errors.Wrap(err, "creating test config directory"))
-		}
-	}
-	if paths.Data != "" {
-		dataDir := filepath.Join(paths.Data, consts.DnoteDirName)
-		if err := os.MkdirAll(dataDir, 0755); err != nil {
-			t.Fatal(errors.Wrap(err, "creating test data directory"))
-		}
-	}
-	if paths.Cache != "" {
-		cacheDir := filepath.Join(paths.Cache, consts.DnoteDirName)
-		if err := os.MkdirAll(cacheDir, 0755); err != nil {
-			t.Fatal(errors.Wrap(err, "creating test cache directory"))
-		}
-	}
-}
 
 // InitTestCtx initializes a test context with an in-memory database
 // and a temporary directory for all paths
 func InitTestCtx(t *testing.T) DnoteCtx {
 	paths := getDefaultTestPaths(t)
 	db := database.InitTestMemoryDB(t)
-	createTestDirectories(t, paths)
+
+	if err := InitDnoteDirs(paths); err != nil {
+		t.Fatal(errors.Wrap(err, "creating test directories"))
+	}
 
 	return DnoteCtx{
 		DB:    db,
@@ -81,7 +62,10 @@ func InitTestCtx(t *testing.T) DnoteCtx {
 // Used when you need full control over database initialization (e.g. migration tests).
 func InitTestCtxWithDB(t *testing.T, db *database.DB) DnoteCtx {
 	paths := getDefaultTestPaths(t)
-	createTestDirectories(t, paths)
+
+	if err := InitDnoteDirs(paths); err != nil {
+		t.Fatal(errors.Wrap(err, "creating test directories"))
+	}
 
 	return DnoteCtx{
 		DB:    db,
@@ -94,7 +78,10 @@ func InitTestCtxWithDB(t *testing.T, db *database.DB) DnoteCtx {
 // at the expected path.
 func InitTestCtxWithFileDB(t *testing.T) DnoteCtx {
 	paths := getDefaultTestPaths(t)
-	createTestDirectories(t, paths)
+
+	if err := InitDnoteDirs(paths); err != nil {
+		t.Fatal(errors.Wrap(err, "creating test directories"))
+	}
 
 	dbPath := filepath.Join(paths.Data, consts.DnoteDirName, consts.DnoteDBFileName)
 	db, err := database.Open(dbPath)
@@ -106,7 +93,6 @@ func InitTestCtxWithFileDB(t *testing.T) DnoteCtx {
 		t.Fatal(errors.Wrap(err, "running schema sql"))
 	}
 
-	database.MarkMigrationComplete(t, db)
 	t.Cleanup(func() { db.Close() })
 
 	return DnoteCtx{
diff --git a/pkg/cli/database/models_test.go b/pkg/cli/database/models_test.go
index 28bc2212..53565d7b 100644
--- a/pkg/cli/database/models_test.go
+++ b/pkg/cli/database/models_test.go
@@ -110,7 +110,6 @@ func TestNoteInsert(t *testing.T) {
 		func() {
 			// Setup
 			db := InitTestMemoryDB(t)
-			defer db.Close()
 
 			n := Note{
 				UUID:     tc.uuid,
@@ -244,7 +243,6 @@ func TestNoteUpdate(t *testing.T) {
 		func() {
 			// Setup
 			db := InitTestMemoryDB(t)
-			defer db.Close()
 
 			n1 := Note{
 				UUID:     tc.uuid,
@@ -267,8 +265,8 @@ func TestNoteUpdate(t *testing.T) {
 				Dirty:    false,
 			}
 
-			MustExec(t, fmt.Sprintf("inserting n1 for test case %d", idx), db, "INSERT INTO notes (uuid, book_uuid, usn, added_on, edited_on, body, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", n1.UUID, n1.BookUUID, n1.USN, n1.AddedOn, n1.EditedOn, n1.Body,  n1.Deleted, n1.Dirty)
-			MustExec(t, fmt.Sprintf("inserting n2 for test case %d", idx), db, "INSERT INTO notes (uuid, book_uuid, usn, added_on, edited_on, body, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", n2.UUID, n2.BookUUID, n2.USN, n2.AddedOn, n2.EditedOn, n2.Body,  n2.Deleted, n2.Dirty)
+			MustExec(t, fmt.Sprintf("inserting n1 for test case %d", idx), db, "INSERT INTO notes (uuid, book_uuid, usn, added_on, edited_on, body, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", n1.UUID, n1.BookUUID, n1.USN, n1.AddedOn, n1.EditedOn, n1.Body, n1.Deleted, n1.Dirty)
+			MustExec(t, fmt.Sprintf("inserting n2 for test case %d", idx), db, "INSERT INTO notes (uuid, book_uuid, usn, added_on, edited_on, body, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", n2.UUID, n2.BookUUID, n2.USN, n2.AddedOn, n2.EditedOn, n2.Body, n2.Deleted, n2.Dirty)
 
 			// execute
 			tx, err := db.Begin()
@@ -336,7 +334,6 @@ func TestNoteUpdateUUID(t *testing.T) {
 		t.Run(fmt.Sprintf("testCase%d", idx), func(t *testing.T) {
 			// Setup
 			db := InitTestMemoryDB(t)
-			defer db.Close()
 
 			n1 := Note{
 				UUID:     "n1-uuid",
@@ -391,7 +388,6 @@ func TestNoteUpdateUUID(t *testing.T) {
 func TestNoteExpunge(t *testing.T) {
 	// Setup
 	db := InitTestMemoryDB(t)
-	defer db.Close()
 
 	n1 := Note{
 		UUID:     "n1-uuid",
@@ -414,8 +410,8 @@ func TestNoteExpunge(t *testing.T) {
 		Dirty:    false,
 	}
 
-	MustExec(t, "inserting n1", db, "INSERT INTO notes (uuid, book_uuid, usn, added_on, edited_on, body, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", n1.UUID, n1.BookUUID, n1.USN, n1.AddedOn, n1.EditedOn, n1.Body,  n1.Deleted, n1.Dirty)
-	MustExec(t, "inserting n2", db, "INSERT INTO notes (uuid, book_uuid, usn, added_on, edited_on, body, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", n2.UUID, n2.BookUUID, n2.USN, n2.AddedOn, n2.EditedOn, n2.Body,  n2.Deleted, n2.Dirty)
+	MustExec(t, "inserting n1", db, "INSERT INTO notes (uuid, book_uuid, usn, added_on, edited_on, body, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", n1.UUID, n1.BookUUID, n1.USN, n1.AddedOn, n1.EditedOn, n1.Body, n1.Deleted, n1.Dirty)
+	MustExec(t, "inserting n2", db, "INSERT INTO notes (uuid, book_uuid, usn, added_on, edited_on, body, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", n2.UUID, n2.BookUUID, n2.USN, n2.AddedOn, n2.EditedOn, n2.Body, n2.Deleted, n2.Dirty)
 
 	// execute
 	tx, err := db.Begin()
@@ -514,7 +510,6 @@ func TestBookInsert(t *testing.T) {
 		func() {
 			// Setup
 			db := InitTestMemoryDB(t)
-			defer db.Close()
 
 			b := Book{
 				UUID:    tc.uuid,
@@ -595,7 +590,6 @@ func TestBookUpdate(t *testing.T) {
 		func() {
 			// Setup
 			db := InitTestMemoryDB(t)
-			defer db.Close()
 
 			b1 := Book{
 				UUID:    "b1-uuid",
@@ -674,7 +668,6 @@ func TestBookUpdateUUID(t *testing.T) {
 
 			// Setup
 			db := InitTestMemoryDB(t)
-			defer db.Close()
 
 			b1 := Book{
 				UUID:    "b1-uuid",
@@ -725,7 +718,6 @@ func TestBookUpdateUUID(t *testing.T) {
 func TestBookExpunge(t *testing.T) {
 	// Setup
 	db := InitTestMemoryDB(t)
-	defer db.Close()
 
 	b1 := Book{
 		UUID:    "b1-uuid",
@@ -780,7 +772,6 @@ func TestBookExpunge(t *testing.T) {
 func TestNoteFTS(t *testing.T) {
 	// set up
 	db := InitTestMemoryDB(t)
-	defer db.Close()
 
 	// execute - insert
 	n := Note{
diff --git a/pkg/cli/database/queries_test.go b/pkg/cli/database/queries_test.go
index d3c2793e..cccc697f 100644
--- a/pkg/cli/database/queries_test.go
+++ b/pkg/cli/database/queries_test.go
@@ -48,7 +48,6 @@ func TestInsertSystem(t *testing.T) {
 		t.Run(fmt.Sprintf("insert %s %s", tc.key, tc.val), func(t *testing.T) {
 			// Setup
 			db := InitTestMemoryDB(t)
-			defer db.Close()
 
 			// execute
 			tx, err := db.Begin()
@@ -96,7 +95,6 @@ func TestUpsertSystem(t *testing.T) {
 		t.Run(fmt.Sprintf("insert %s %s", tc.key, tc.val), func(t *testing.T) {
 			// Setup
 			db := InitTestMemoryDB(t)
-			defer db.Close()
 
 			MustExec(t, "inserting a system configuration", db, "INSERT INTO system (key, value) VALUES (?, ?)", "baz", "quz")
 
@@ -135,7 +133,6 @@ func TestGetSystem(t *testing.T) {
 	t.Run(fmt.Sprintf("get string value"), func(t *testing.T) {
 		// Setup
 		db := InitTestMemoryDB(t)
-		defer db.Close()
 
 		// execute
 		MustExec(t, "inserting a system configuration", db, "INSERT INTO system (key, value) VALUES (?, ?)", "foo", "bar")
@@ -158,7 +155,6 @@ func TestGetSystem(t *testing.T) {
 	t.Run(fmt.Sprintf("get int64 value"), func(t *testing.T) {
 		// Setup
 		db := InitTestMemoryDB(t)
-		defer db.Close()
 
 		// execute
 		MustExec(t, "inserting a system configuration", db, "INSERT INTO system (key, value) VALUES (?, ?)", "foo", 1234)
@@ -199,7 +195,6 @@ func TestUpdateSystem(t *testing.T) {
 		t.Run(fmt.Sprintf("update %s %s", tc.key, tc.val), func(t *testing.T) {
 			// Setup
 			db := InitTestMemoryDB(t)
-			defer db.Close()
 
 			MustExec(t, "inserting a system configuration", db, "INSERT INTO system (key, value) VALUES (?, ?)", "foo", "fuz")
 			MustExec(t, "inserting a system configuration", db, "INSERT INTO system (key, value) VALUES (?, ?)", "baz", "quz")
@@ -239,7 +234,6 @@ func TestGetActiveNote(t *testing.T) {
 	t.Run("not deleted", func(t *testing.T) {
 		// set up
 		db := InitTestMemoryDB(t)
-		defer db.Close()
 
 		n1UUID := "n1-uuid"
 		MustExec(t, "inserting n1", db, "INSERT INTO notes (uuid, book_uuid, body, added_on, edited_on, usn, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", n1UUID, "b1-uuid", "n1 content", 1542058875, 1542058876, 1,  false, true)
@@ -268,7 +262,6 @@ func TestGetActiveNote(t *testing.T) {
 	t.Run("deleted", func(t *testing.T) {
 		// set up
 		db := InitTestMemoryDB(t)
-		defer db.Close()
 
 		n1UUID := "n1-uuid"
 		MustExec(t, "inserting n1", db, "INSERT INTO notes (uuid, book_uuid, body, added_on, edited_on, usn, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", n1UUID, "b1-uuid", "n1 content", 1542058875, 1542058876, 1,  true, true)
@@ -292,7 +285,6 @@ func TestGetActiveNote(t *testing.T) {
 func TestUpdateNoteContent(t *testing.T) {
 	// set up
 	db := InitTestMemoryDB(t)
-	defer db.Close()
 
 	uuid := "n1-uuid"
 	MustExec(t, "inserting n1", db, "INSERT INTO notes (uuid, book_uuid, body, added_on, edited_on, usn, deleted, dirty) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", uuid, "b1-uuid", "n1 content", 1542058875, 0, 1,  false, false)
@@ -324,7 +316,6 @@ func TestUpdateNoteContent(t *testing.T) {
 func TestUpdateNoteBook(t *testing.T) {
 	// set up
 	db := InitTestMemoryDB(t)
-	defer db.Close()
 
 	b1UUID := "b1-uuid"
 	b2UUID := "b2-uuid"
@@ -361,7 +352,6 @@ func TestUpdateNoteBook(t *testing.T) {
 func TestUpdateBookName(t *testing.T) {
 	// set up
 	db := InitTestMemoryDB(t)
-	defer db.Close()
 
 	b1UUID := "b1-uuid"
 	MustExec(t, "inserting b1", db, "INSERT INTO books (uuid, label, usn, deleted, dirty) VALUES (?, ?, ?, ?, ?)", b1UUID, "b1-label", 8, false, false)
diff --git a/pkg/cli/database/schema.sql b/pkg/cli/database/schema.sql
index 655148c8..9a9de094 100644
--- a/pkg/cli/database/schema.sql
+++ b/pkg/cli/database/schema.sql
@@ -34,3 +34,7 @@ CREATE TRIGGER notes_after_update AFTER UPDATE ON notes BEGIN
 				INSERT INTO note_fts(note_fts, rowid, body) VALUES ('delete', old.rowid, old.body);
 				INSERT INTO note_fts(rowid, body) VALUES (new.rowid, new.body);
 			END;
+
+-- Migration version data.
+INSERT INTO system (key, value) VALUES ('schema', 14);
+INSERT INTO system (key, value) VALUES ('remote_schema', 1);
diff --git a/pkg/cli/database/schema/main.go b/pkg/cli/database/schema/main.go
index e4e219de..2596c84a 100644
--- a/pkg/cli/database/schema/main.go
+++ b/pkg/cli/database/schema/main.go
@@ -26,6 +26,7 @@ import (
 	"strings"
 
 	"github.com/dnote/dnote/pkg/cli/config"
+	"github.com/dnote/dnote/pkg/cli/consts"
 	"github.com/dnote/dnote/pkg/cli/context"
 	"github.com/dnote/dnote/pkg/cli/database"
 	"github.com/dnote/dnote/pkg/cli/infra"
@@ -118,7 +119,12 @@ func generateSchema(tmpDir string) (string, error) {
 		return "", fmt.Errorf("extracting schema: %w", err)
 	}
 
-	return schema, nil
+	// Add INSERT statements for migration versions.
+	systemData := "\n-- Migration version data.\n"
+	systemData += fmt.Sprintf("INSERT INTO system (key, value) VALUES ('%s', %d);\n", consts.SystemSchema, len(migrate.LocalSequence))
+	systemData += fmt.Sprintf("INSERT INTO system (key, value) VALUES ('%s', %d);\n", consts.SystemRemoteSchema, len(migrate.RemoteSequence))
+
+	return schema + systemData, nil
 }
 
 // extractSchema extracts the complete schema by querying sqlite_master
diff --git a/pkg/cli/database/schema/main_test.go b/pkg/cli/database/schema/main_test.go
index acc46654..24bd9fc0 100644
--- a/pkg/cli/database/schema/main_test.go
+++ b/pkg/cli/database/schema/main_test.go
@@ -26,6 +26,7 @@ import (
 	"testing"
 
 	"github.com/dnote/dnote/pkg/assert"
+	"github.com/dnote/dnote/pkg/cli/consts"
 )
 
 func TestRun(t *testing.T) {
@@ -73,4 +74,11 @@ func TestRun(t *testing.T) {
 
 	// Verify schema does not contain sqlite internal tables
 	assert.Equal(t, strings.Contains(schema, "sqlite_sequence"), false, "schema should not contain sqlite_sequence")
+
+	// Verify system key-value pairs for schema versions are present
+	expectedSchemaKey := fmt.Sprintf("INSERT INTO system (key, value) VALUES ('%s',", consts.SystemSchema)
+	assert.Equal(t, strings.Contains(schema, expectedSchemaKey), true, "schema should contain schema version INSERT statement")
+
+	expectedRemoteSchemaKey := fmt.Sprintf("INSERT INTO system (key, value) VALUES ('%s',", consts.SystemRemoteSchema)
+	assert.Equal(t, strings.Contains(schema, expectedRemoteSchemaKey), true, "schema should contain remote_schema version INSERT statement")
 }
diff --git a/pkg/cli/database/test_embed.go b/pkg/cli/database/test_embed.go
deleted file mode 100644
index c2fe08e9..00000000
--- a/pkg/cli/database/test_embed.go
+++ /dev/null
@@ -1,27 +0,0 @@
-package database
-
-import (
-	"testing"
-)
-
-func TestSchemaEmbed(t *testing.T) {
-	db := InitTestMemoryDB(t)
-	defer db.Close()
-
-	// Try to insert a book to verify schema is loaded
-	_, err := db.Exec("INSERT INTO books (uuid, label) VALUES (?, ?)", "test-uuid", "test-label")
-	if err != nil {
-		t.Fatalf("Failed to insert into books: %v", err)
-	}
-
-	// Verify it was inserted
-	var label string
-	err = db.QueryRow("SELECT label FROM books WHERE uuid = ?", "test-uuid").Scan(&label)
-	if err != nil {
-		t.Fatalf("Failed to query books: %v", err)
-	}
-	if label != "test-label" {
-		t.Fatalf("Expected label 'test-label', got '%s'", label)
-	}
-	t.Log("Schema embed test passed!")
-}
diff --git a/pkg/cli/database/testutils.go b/pkg/cli/database/testutils.go
index 602a48ba..16685686 100644
--- a/pkg/cli/database/testutils.go
+++ b/pkg/cli/database/testutils.go
@@ -58,9 +58,7 @@ func MustExec(t *testing.T, message string, db *DB, query string, args ...interf
 
 // InitTestMemoryDB initializes an in-memory test database with the default schema.
 func InitTestMemoryDB(t *testing.T) *DB {
-	db := InitTestMemoryDBRaw(t, "")
-	MarkMigrationComplete(t, db)
-	return db
+	return InitTestMemoryDBRaw(t, "")
 }
 
 // InitTestFileDB initializes a file-based test database with the default schema.
@@ -81,8 +79,6 @@ func InitTestFileDBRaw(t *testing.T, dbPath string) *DB {
 		t.Fatal(errors.Wrap(err, "running schema sql"))
 	}
 
-	MarkMigrationComplete(t, db)
-
 	t.Cleanup(func() { db.Close() })
 	return db
 }
@@ -125,16 +121,6 @@ func OpenTestDB(t *testing.T, dnoteDir string) *DB {
 	return db
 }
 
-// MarkMigrationComplete marks all migrations as complete in the database
-func MarkMigrationComplete(t *testing.T, db *DB) {
-	if _, err := db.Exec("INSERT INTO system (key, value) VALUES (? , ?);", consts.SystemSchema, 14); err != nil {
-		t.Fatal(errors.Wrap(err, "inserting schema"))
-	}
-	if _, err := db.Exec("INSERT INTO system (key, value) VALUES (? , ?);", consts.SystemRemoteSchema, 1); err != nil {
-		t.Fatal(errors.Wrap(err, "inserting remote schema"))
-	}
-}
-
 // mustGenerateTestUUID generates a UUID for test databases and fails the test on error
 func mustGenerateTestUUID(t *testing.T) string {
 	uuid, err := utils.GenerateUUID()
diff --git a/pkg/cli/infra/init.go b/pkg/cli/infra/init.go
index fbff4267..1e572cf7 100644
--- a/pkg/cli/infra/init.go
+++ b/pkg/cli/infra/init.go
@@ -24,7 +24,6 @@ import (
 	"database/sql"
 	"fmt"
 	"os"
-	"path/filepath"
 	"strconv"
 	"time"
 
@@ -329,37 +328,6 @@ func getEditorCommand() string {
 	return ret
 }
 
-// initDir creates a directory if it doesn't exist
-func initDir(path string) error {
-	ok, err := utils.FileExists(path)
-	if err != nil {
-		return errors.Wrapf(err, "checking if dir exists at %s", path)
-	}
-	if ok {
-		return nil
-	}
-
-	if err := os.MkdirAll(path, 0755); err != nil {
-		return errors.Wrapf(err, "creating a directory at %s", path)
-	}
-
-	return nil
-}
-
-// initDnoteDir initializes missing directories that Dnote uses
-func initDnoteDir(ctx context.DnoteCtx) error {
-	if err := initDir(filepath.Join(ctx.Paths.Config, consts.DnoteDirName)); err != nil {
-		return errors.Wrap(err, "initializing config dir")
-	}
-	if err := initDir(filepath.Join(ctx.Paths.Data, consts.DnoteDirName)); err != nil {
-		return errors.Wrap(err, "initializing data dir")
-	}
-	if err := initDir(filepath.Join(ctx.Paths.Cache, consts.DnoteDirName)); err != nil {
-		return errors.Wrap(err, "initializing cache dir")
-	}
-
-	return nil
-}
 
 // initConfigFile populates a new config file if it does not exist yet
 func initConfigFile(ctx context.DnoteCtx, apiEndpoint string) error {
@@ -395,7 +363,7 @@ func initConfigFile(ctx context.DnoteCtx, apiEndpoint string) error {
 
 // initFiles creates, if necessary, the dnote directory and files inside
 func initFiles(ctx context.DnoteCtx, apiEndpoint string) error {
-	if err := initDnoteDir(ctx); err != nil {
+	if err := context.InitDnoteDirs(ctx.Paths); err != nil {
 		return errors.Wrap(err, "creating the dnote dir")
 	}
 	if err := initConfigFile(ctx, apiEndpoint); err != nil {
diff --git a/pkg/cli/infra/init_test.go b/pkg/cli/infra/init_test.go
index 401294d5..6ac0d5cf 100644
--- a/pkg/cli/infra/init_test.go
+++ b/pkg/cli/infra/init_test.go
@@ -33,7 +33,6 @@ import (
 func TestInitSystemKV(t *testing.T) {
 	// Setup
 	db := database.InitTestMemoryDB(t)
-	defer db.Close()
 
 	var originalCount int
 	database.MustScan(t, "counting system configs", db.QueryRow("SELECT count(*) FROM system"), &originalCount)
@@ -65,7 +64,6 @@ func TestInitSystemKV(t *testing.T) {
 func TestInitSystemKV_existing(t *testing.T) {
 	// Setup
 	db := database.InitTestMemoryDB(t)
-	defer db.Close()
 
 	database.MustExec(t, "inserting a system config", db, "INSERT INTO system (key, value) VALUES (?, ?)", "testKey", "testVal")
 
diff --git a/pkg/cli/migrate/migrate_test.go b/pkg/cli/migrate/migrate_test.go
index 4bb66057..591f1800 100644
--- a/pkg/cli/migrate/migrate_test.go
+++ b/pkg/cli/migrate/migrate_test.go
@@ -38,6 +38,14 @@ import (
 	"github.com/pkg/errors"
 )
 
+// initTestDBNoMigration initializes a test database with schema.sql but removes
+// migration version data so tests can control the migration state themselves.
+func initTestDBNoMigration(t *testing.T) *database.DB {
+	db := database.InitTestMemoryDBRaw(t, "")
+	// Remove migration versions from schema.sql so tests can set their own
+	database.MustExec(t, "clearing schema versions", db, "DELETE FROM system WHERE key IN (?, ?)", consts.SystemSchema, consts.SystemRemoteSchema)
+	return db
+}
 
 func TestExecute_bump_schema(t *testing.T) {
 	testCases := []struct {
@@ -54,7 +62,7 @@ func TestExecute_bump_schema(t *testing.T) {
 	for _, tc := range testCases {
 		func() {
 			// set up
-			db := database.InitTestMemoryDBRaw(t, "")
+			db := initTestDBNoMigration(t)
 			ctx := context.InitTestCtxWithDB(t, db)
 
 			database.MustExec(t, "inserting a schema", db, "INSERT INTO system (key, value) VALUES (?, ?)", tc.schemaKey, 8)
@@ -108,7 +116,7 @@ func TestRun_nonfresh(t *testing.T) {
 	for _, tc := range testCases {
 		func() {
 			// set up
-			db := database.InitTestMemoryDBRaw(t, "")
+			db := initTestDBNoMigration(t)
 			ctx := context.InitTestCtxWithDB(t, db)
 			database.MustExec(t, "inserting a schema", db, "INSERT INTO system (key, value) VALUES (?, ?)", tc.schemaKey, 2)
 			database.MustExec(t, "creating a temporary table for testing", db,
@@ -185,7 +193,7 @@ func TestRun_fresh(t *testing.T) {
 	for _, tc := range testCases {
 		func() {
 			// set up
-			db := database.InitTestMemoryDBRaw(t, "")
+			db := initTestDBNoMigration(t)
 			ctx := context.InitTestCtxWithDB(t, db)
 
 			database.MustExec(t, "creating a temporary table for testing", db,
@@ -256,7 +264,7 @@ func TestRun_up_to_date(t *testing.T) {
 	for _, tc := range testCases {
 		func() {
 			// set up
-			db := database.InitTestMemoryDBRaw(t, "")
+			db := initTestDBNoMigration(t)
 			ctx := context.InitTestCtxWithDB(t, db)
 
 			database.MustExec(t, "creating a temporary table for testing", db,
diff --git a/pkg/cli/utils/files.go b/pkg/cli/utils/files.go
index b4b2d1df..00e898ec 100644
--- a/pkg/cli/utils/files.go
+++ b/pkg/cli/utils/files.go
@@ -55,6 +55,24 @@ func FileExists(filepath string) (bool, error) {
 	return false, errors.Wrap(err, "getting file info")
 }
 
+// EnsureDir creates a directory if it doesn't exist.
+// Returns nil if the directory already exists or was successfully created.
+func EnsureDir(path string) error {
+	ok, err := FileExists(path)
+	if err != nil {
+		return errors.Wrapf(err, "checking if dir exists at %s", path)
+	}
+	if ok {
+		return nil
+	}
+
+	if err := os.MkdirAll(path, 0755); err != nil {
+		return errors.Wrapf(err, "creating directory at %s", path)
+	}
+
+	return nil
+}
+
 // CopyDir copies a directory from src to dest, recursively copying nested
 // directories
 func CopyDir(src, dest string) error {
diff --git a/pkg/cli/utils/files_test.go b/pkg/cli/utils/files_test.go
new file mode 100644
index 00000000..3c7c92bd
--- /dev/null
+++ b/pkg/cli/utils/files_test.go
@@ -0,0 +1,45 @@
+/* Copyright (C) 2019, 2020, 2021, 2022, 2023, 2024, 2025 Dnote contributors
+ *
+ * This file is part of Dnote.
+ *
+ * Dnote is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * Dnote is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Dnote.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package utils
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/dnote/dnote/pkg/assert"
+)
+
+func TestEnsureDir(t *testing.T) {
+	tmpDir := t.TempDir()
+	testPath := filepath.Join(tmpDir, "test", "nested", "dir")
+
+	// Create directory
+	err := EnsureDir(testPath)
+	assert.Equal(t, err, nil, "EnsureDir should succeed")
+
+	// Verify it exists
+	info, err := os.Stat(testPath)
+	assert.Equal(t, err, nil, "directory should exist")
+	assert.Equal(t, info.IsDir(), true, "should be a directory")
+
+	// Call again on existing directory - should not error
+	err = EnsureDir(testPath)
+	assert.Equal(t, err, nil, "EnsureDir should succeed on existing directory")
+}