Add State Restoration Tests, DB Recovery Steps

Author: thecoolwinter

CodeEdit/Features/Editor/Models/Restoration/EditorStateRestoration.swift

@@ -61,20 +61,25 @@ final class EditorStateRestoration {
         }
     }
 
-    private var databaseQueue: DatabaseQueue
-
-    private init() throws {
-        let databaseURL: URL = FileManager.default
+    private var databaseQueue: DatabaseQueue?
+    private var databaseURL: URL
+
+    /// Create a new editor restoration object. Will connect to or create a SQLite db.
+    /// - Parameter databaseURL: The database URL to use. Must point to a file, not a directory. If left `nil`, will
+    ///                          create a new database named `editor-restoration.db` in the application support
+    ///                          directory.
+    init(_ databaseURL: URL? = nil) throws {
+        self.databaseURL = databaseURL ?? FileManager.default
             .homeDirectoryForCurrentUser
             .appending(path: "Library/Application Support/CodeEdit", directoryHint: .isDirectory)
             .appending(path: "editor-restoration.db", directoryHint: .notDirectory)
-
-        self.databaseQueue = try DatabaseQueue(path: databaseURL.absolutePath, configuration: .init())
         try attemptMigration(retry: true)
     }
 
-    private func attemptMigration(retry: Bool) throws {
+    func attemptMigration(retry: Bool) throws {
         do {
+            let databaseQueue = try DatabaseQueue(path: self.databaseURL.absolutePath, configuration: .init())
+
             var migrator = DatabaseMigrator()
 
             migrator.registerMigration("Version 0") {
@@ -85,12 +90,11 @@ final class EditorStateRestoration {
             }
 
             try migrator.migrate(databaseQueue)
+            self.databaseQueue = databaseQueue
         } catch {
             if retry {
                 // Try to delete the database on failure, might fix a corruption or version error.
                 try? FileManager.default.removeItem(at: databaseURL)
-                // This will recreate the db file if necessary
-                self.databaseQueue = try DatabaseQueue(path: databaseURL.absolutePath, configuration: .init())
                 try attemptMigration(retry: false)
 
                 return // Ignore the original error if we're retrying
@@ -108,7 +112,7 @@ final class EditorStateRestoration {
         do {
             let serializedData = try JSONEncoder().encode(data)
             let dbRow = StateRestorationRecord(uri: documentUrl.absolutePath, data: serializedData)
-            try databaseQueue.write { try dbRow.upsert($0) }
+            try databaseQueue?.write { try dbRow.upsert($0) }
         } catch {
             Self.logger.error("Failed to save editor state: \(error)")
         }
@@ -119,7 +123,7 @@ final class EditorStateRestoration {
     /// - Returns: Any data saved for this file.
     func restorationState(for documentUrl: URL) -> StateRestorationData? {
         do {
-            guard let row = try databaseQueue.read({
+            guard let row = try databaseQueue?.read({
                 try StateRestorationRecord.fetchOne($0, key: documentUrl.absolutePath)
             }) else {
                 return nil

CodeEdit/Features/Editor/Models/Restoration/UndoManagerRegistration.swift

@@ -19,7 +19,7 @@ final class UndoManagerRegistration: ObservableObject {
     private var managerMap: [CEWorkspaceFile.ID: CEUndoManager] = [:]
 
     init() { }
-    
+
     /// Find or create a new undo manager.
     /// - Parameter file: The file to create for.
     /// - Returns: The undo manager for the given file.

CodeEdit/Features/Editor/Views/CodeFileView.swift

@@ -58,8 +58,7 @@ struct CodeFileView: View {
     @Environment(\.colorScheme)
     private var colorScheme
 
-    @EnvironmentObject
-    var undoRegistration: UndoManagerRegistration
+    @EnvironmentObject var undoRegistration: UndoManagerRegistration
 
     @ObservedObject private var themeModel: ThemeModel = .shared
 

CodeEditTests/Features/Editor/EditorStateRestorationTests.swift

@@ -0,0 +1,49 @@
+//
+//  EditorStateRestorationTests.swift
+//  CodeEditTests
+//
+//  Created by Khan Winter on 7/3/25.
+//
+
+import Testing
+import Foundation
+@testable import CodeEdit
+
+@Suite(.serialized)
+struct EditorStateRestorationTests {
+    @Test
+    func createsDatabase() throws {
+        try withTempDir { dir in
+            let url = dir.appending(path: "database.db")
+            _ = try EditorStateRestoration(url)
+            #expect(FileManager.default.fileExists(atPath: url.path(percentEncoded: false)))
+        }
+    }
+
+    @Test
+    func savesAndRetrievesState() throws {
+        try withTempDir { dir in
+            let url = dir.appending(path: "database.db")
+            let restoration = try EditorStateRestoration(url)
+
+            // Update some state
+            restoration.updateRestorationState(
+                for: dir.appending(path: "file.txt"),
+                data: .init(cursorPositions: [], scrollPosition: .zero)
+            )
+
+            // Retrieve it
+            #expect(restoration.restorationState(for: dir.appending(path: "file.txt")) != nil)
+        }
+    }
+
+    @Test
+    func clearsCorruptedDatabase() throws {
+        try withTempDir { dir in
+            let url = dir.appending(path: "database.db")
+            try "bad SQLITE data HAHAHA".write(to: url, atomically: true, encoding: .utf8)
+            // This will throw if it can't connect to the database.
+            _ = try EditorStateRestoration(url)
+        }
+    }
+}

CodeEditTests/Utils/withTempDir.swift

@@ -0,0 +1,53 @@
+//
+//  withTempDir.swift
+//  CodeEditTests
+//
+//  Created by Khan Winter on 7/3/25.
+//
+
+import Foundation
+
+func withTempDir(_ test: (URL) async throws -> Void) async throws {
+    let tempDirURL = try createAndClearDir()
+    do {
+        try await test(tempDirURL)
+    } catch {
+        try clearDir(tempDirURL)
+        throw error
+    }
+    try clearDir(tempDirURL)
+}
+
+func withTempDir(_ test: (URL) throws -> Void) throws {
+    let tempDirURL = try createAndClearDir()
+    do {
+        try test(tempDirURL)
+    } catch {
+        try clearDir(tempDirURL)
+        throw error
+    }
+    try clearDir(tempDirURL)
+}
+
+private func createAndClearDir() throws -> URL {
+    let tempDirURL = try FileManager.default.url(
+        for: .developerApplicationDirectory,
+        in: .userDomainMask,
+        appropriateFor: nil,
+        create: true
+    )
+        .appending(path: "CodeEditTestDirectory", directoryHint: .isDirectory)
+
+    // If it exists, delete it before the test
+    try clearDir(tempDirURL)
+
+    try FileManager.default.createDirectory(at: tempDirURL, withIntermediateDirectories: true)
+
+    return tempDirURL
+}
+
+private func clearDir(_ url: URL) throws {
+    if FileManager.default.fileExists(atPath: url.absoluteURL.path(percentEncoded: false)) {
+        try FileManager.default.removeItem(at: url)
+    }
+}