diff --git a/src/display/annotation_storage.js b/src/display/annotation_storage.js
index 843e8cd74..4fe1f478a 100644
--- a/src/display/annotation_storage.js
+++ b/src/display/annotation_storage.js
@@ -120,10 +120,22 @@ class AnnotationStorage {
     return this.#storage.has(key);
   }
 
+  /**
+   * @returns {Object | null}
+   */
   getAll() {
     return this.#storage.size > 0 ? objectFromMap(this.#storage) : null;
   }
 
+  /**
+   * @param {Object} obj
+   */
+  setAll(obj) {
+    for (const [key, val] of Object.entries(obj)) {
+      this.setValue(key, val);
+    }
+  }
+
   get size() {
     return this.#storage.size;
   }
diff --git a/test/driver.js b/test/driver.js
index fc6f092ce..410173c9d 100644
--- a/test/driver.js
+++ b/test/driver.js
@@ -486,24 +486,18 @@ class Driver {
         let promise = loadingTask.promise;
 
         if (task.save) {
-          if (!task.annotationStorage) {
-            promise = Promise.reject(
-              new Error("Missing `annotationStorage` entry.")
-            );
-          } else {
-            promise = loadingTask.promise.then(async doc => {
-              for (const [key, value] of Object.entries(
-                task.annotationStorage
-              )) {
-                doc.annotationStorage.setValue(key, value);
-              }
-              const data = await doc.saveDocument();
-              await loadingTask.destroy();
-              delete task.annotationStorage;
+          promise = loadingTask.promise.then(async doc => {
+            if (!task.annotationStorage) {
+              throw new Error("Missing `annotationStorage` entry.");
+            }
+            doc.annotationStorage.setAll(task.annotationStorage);
 
-              return getDocument(data).promise;
-            });
-          }
+            const data = await doc.saveDocument();
+            await loadingTask.destroy();
+            delete task.annotationStorage;
+
+            return getDocument(data).promise;
+          });
         }
 
         promise.then(
@@ -668,11 +662,7 @@ class Driver {
               pageColors = null;
 
             if (task.annotationStorage) {
-              const entries = Object.entries(task.annotationStorage),
-                docAnnotationStorage = task.pdfDoc.annotationStorage;
-              for (const [key, value] of entries) {
-                docAnnotationStorage.setValue(key, value);
-              }
+              task.pdfDoc.annotationStorage.setAll(task.annotationStorage);
             }
 
             let textLayerCanvas, annotationLayerCanvas, annotationLayerContext;