diff --git a/MediaEditor.podspec b/MediaEditor.podspec
index 753dc38..cfa444b 100644
--- a/MediaEditor.podspec
+++ b/MediaEditor.podspec
@@ -1,6 +1,6 @@
 Pod::Spec.new do |s|
   s.name             = 'MediaEditor'
-  s.version          = '1.1.0-beta.1'
+  s.version          = '1.1.0-beta.2'
   s.summary          = 'An extensible Media Editor for iOS.'
 
   s.description      = <<-DESC
diff --git a/Sources/Capabilities/Drawing/MediaEditorDrawing.swift b/Sources/Capabilities/Drawing/MediaEditorDrawing.swift
index 9ffdcfb..dd5d643 100644
--- a/Sources/Capabilities/Drawing/MediaEditorDrawing.swift
+++ b/Sources/Capabilities/Drawing/MediaEditorDrawing.swift
@@ -55,12 +55,13 @@ class MediaEditorDrawing: UIViewController {
     }
 
     @IBAction func done(_ sender: Any) {
-        guard let image = annotationView.image else {
+        guard annotationView.canUndo,
+            let image = annotationView.image else {
             onCancel?()
             return
         }
 
-        onFinishEditing?(image, [.other])
+        onFinishEditing?(image, [.draw])
     }
 }
 
diff --git a/Sources/Enums/MediaEditorOperation.swift b/Sources/Enums/MediaEditorOperation.swift
index 9483fac..a4d75ed 100644
--- a/Sources/Enums/MediaEditorOperation.swift
+++ b/Sources/Enums/MediaEditorOperation.swift
@@ -5,4 +5,5 @@ public enum MediaEditorOperation {
     case rotate
     case filter
     case other
+    case draw
 }
diff --git a/Tests/Capabilities/Drawing/MediaEditorDrawingTests.swift b/Tests/Capabilities/Drawing/MediaEditorDrawingTests.swift
index b9c325c..2370617 100644
--- a/Tests/Capabilities/Drawing/MediaEditorDrawingTests.swift
+++ b/Tests/Capabilities/Drawing/MediaEditorDrawingTests.swift
@@ -60,8 +60,10 @@ class MediaEditorDrawingTests: XCTestCase {
         let mediaEditorDrawing = MediaEditorDrawing.initialize(image, onFinishEditing: { finishedImage, _ in
             result = finishedImage
         }, onCancel: {}) as! MediaEditorDrawing
+        let annotationViewMock = MediaEditorAnnotationViewMock()
 
         mediaEditorDrawing.loadView()
+        mediaEditorDrawing.annotationView = annotationViewMock
         mediaEditorDrawing.viewDidLoad()
         mediaEditorDrawing.done(self)
 
@@ -70,17 +72,17 @@ class MediaEditorDrawingTests: XCTestCase {
 
     func testModifiedImageIsReturnedIfChangesAreMade() {
         let image = UIImage(systemName: "arrowshape.turn.up.left")!
-        let drawingUrl = Bundle(for: MediaEditorDrawingTests.self).url(forResource: "demo-drawing", withExtension: nil)!
-        let drawingData = try! Data(contentsOf: drawingUrl)
 
         var result: UIImage? = nil
         let mediaEditorDrawing = MediaEditorDrawing.initialize(image, onFinishEditing: { finishedImage, _ in
             result = finishedImage
         }, onCancel: {}) as! MediaEditorDrawing
+        let annotationViewMock = MediaEditorAnnotationViewMock()
+        annotationViewMock.image = UIImage()
 
         mediaEditorDrawing.loadView()
         mediaEditorDrawing.viewDidLoad()
-        mediaEditorDrawing.annotationView.drawingData = drawingData
+        mediaEditorDrawing.annotationView = annotationViewMock
         mediaEditorDrawing.view.setNeedsLayout()
         mediaEditorDrawing.view.layoutIfNeeded()
 
@@ -98,3 +100,10 @@ class MediaEditorDrawingTests: XCTestCase {
         expect(mediaEditorDrawing.redoButton.isEnabled).to(beFalse())
     }
 }
+
+@available(iOS 13.0, *)
+private class MediaEditorAnnotationViewMock: MediaEditorAnnotationView {
+    override var canUndo: Bool {
+        return true
+    }
+}