Reset containedEditor when attachment is removed from container (#274)

rajdeep · Feb 9, 2024 · f04eb43 · f04eb43
1 parent 14c352f
commit f04eb43
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/Proton/Sources/Swift/Attachment/Attachment.swift b/Proton/Sources/Swift/Attachment/Attachment.swift
@@ -355,6 +355,7 @@ open class Attachment: NSTextAttachment, BoundsObserving {
     @objc
     func removeFromSuperview() {
         view?.removeFromSuperview()
+        containerEditorView = nil
     }
 
     /// Causes invalidation of layout of the attachment when the containing view bounds are changed

diff --git a/Proton/Tests/Editor/EditorViewTests.swift b/Proton/Tests/Editor/EditorViewTests.swift
@@ -831,6 +831,28 @@ class EditorViewTests: XCTestCase {
         _ = editor.attachmentsInRange(NSRange(location: 2, length: 1))
         XCTAssertEqual(editor.attributedText.length, 2)
     }
+
+    func testRemovesContainerEditorOnDeletingAttachment() {
+        let viewController = EditorTestViewController()
+        let editor = viewController.editor
+
+        let dummyAttachment1 = DummyMultiEditorAttachment(numberOfEditors: 1)
+        editor.insertAttachment(in: .zero, attachment: dummyAttachment1)
+
+        let dummyAttachment2 = DummyMultiEditorAttachment(numberOfEditors: 2)
+        editor.insertAttachment(in: editor.textEndRange, attachment: dummyAttachment2)
+
+        viewController.render()
+
+        XCTAssertEqual(dummyAttachment1.containerEditorView, editor)
+        XCTAssertEqual(dummyAttachment2.containerEditorView, editor)
+
+        editor.replaceCharacters(in: NSRange(location: 0, length: 1), with: NSAttributedString())
+        dummyAttachment2.removeFromContainer()
+
+        XCTAssertNil(dummyAttachment1.containerEditorView)
+        XCTAssertNil(dummyAttachment2.containerEditorView)
+    }
 }
 
 class DummyMultiEditorAttachment: Attachment {