[builder] Add null-safe wrapper for Firestore documents

This will improve the migration of the builder scripts to null-safety.

Change-Id: I9f4c82f637b89609e4963ebe56d3a256c04bd32e
Reviewed-on: https://dart-review.googlesource.com/c/dart_ci/+/223120
Reviewed-by: Alexander Thomas <athom@google.com>
Commit-Queue: William Hesse <whesse@google.com>
diff --git a/builder/lib/src/builder.dart b/builder/lib/src/builder.dart
index 257e5d3..92b8307 100644
--- a/builder/lib/src/builder.dart
+++ b/builder/lib/src/builder.dart
@@ -115,7 +115,7 @@
       if (review != null) {
         tryApprovals.addAll({
           for (final result in await firestore.tryApprovals(review))
-            testResult(result): index
+            testResult(result.fields): index
         });
       }
       if (reverted != null) {
diff --git a/builder/lib/src/firestore.dart b/builder/lib/src/firestore.dart
index 26d9247..91f6435 100644
--- a/builder/lib/src/firestore.dart
+++ b/builder/lib/src/firestore.dart
@@ -15,17 +15,17 @@
 export 'firestore_helpers.dart';
 
 class Commit {
-  final DataWrapper wrapper;
+  final SafeDocument document;
   final String hash;
-  Commit(this.hash, Document document) : wrapper = DataWrapper(document);
+  Commit(this.hash, this.document);
   Commit.fromJson(this.hash, Map<String, dynamic> data)
-      : wrapper = DataWrapper.fields(taggedMap(data));
-  int get index => wrapper.getInt('index');
-  String get revertOf => wrapper.getString(fRevertOf);
-  bool get isRevert => wrapper.fields.containsKey(fRevertOf);
-  int get review => wrapper.getInt(fReview);
+      : document = SafeDocument(Document(fields: taggedMap(data), name: ''));
+  int get index => document.getInt('index');
+  String get revertOf => document.getString(fRevertOf);
+  bool get isRevert => document.fields.containsKey(fRevertOf);
+  int get review => document.getInt(fReview);
 
-  Map<String, Object> toJson() => untagMap(wrapper.fields);
+  Map<String, Object> toJson() => untagMap(document.fields);
 }
 
 class FirestoreService {
@@ -49,7 +49,7 @@
     }
   }
 
-  Future<RunQueryResponse> query(
+  Future<List<SafeDocument>> query(
       {String from,
       Filter where,
       Order orderBy,
@@ -68,14 +68,7 @@
     if (limit != null) {
       query.limit = limit;
     }
-    final responses = await runQuery(query, parent: parent);
-    // The REST API will respond with a single `null` result in case the query
-    // did not match any documents.
-    if (responses.length == 1 && responses.single.document == null) {
-      responses.length = 0;
-    }
-    documentsFetched += responses.length;
-    return responses;
+    return runQuery(query, parent: parent);
   }
 
   Future<Document> getDocument(String path,
@@ -101,10 +94,18 @@
   String get database => 'projects/$project/databases/(default)';
   String get documents => '$database/documents';
 
-  Future<RunQueryResponse> runQuery(StructuredQuery query, {String parent}) {
+  Future<List<SafeDocument>> runQuery(StructuredQuery query,
+      {String parent}) async {
     final request = RunQueryRequest()..structuredQuery = query;
     final parentPath = parent == null ? documents : '$documents/$parent';
-    return firestore.projects.databases.documents.runQuery(request, parentPath);
+    final queryResponse = await firestore.projects.databases.documents
+        .runQuery(request, parentPath);
+    if (queryResponse.first.document == null) return [];
+    documentsFetched += queryResponse.length;
+    return [
+      for (final responseElement in queryResponse)
+        SafeDocument(responseElement.document /*!*/)
+    ];
   }
 
   Future<void> writeDocument(Document document) async {
@@ -119,7 +120,7 @@
     return documentExists('$documents/reviews/$review/patchsets/$patchset');
   }
 
-  Commit _commit(Document document) {
+  Commit _commit(SafeDocument document) {
     document.fields['hash'] = taggedValue(document.name.split('/').last);
     return Commit(document.name.split('/').last, document);
   }
@@ -127,19 +128,19 @@
   Future<Commit> getCommit(String hash) async {
     final document =
         await getDocument('$documents/commits/$hash', throwOnNotFound: false);
-    return document != null ? Commit(hash, document) : null;
+    return document != null ? Commit(hash, SafeDocument(document)) : null;
   }
 
   Future<Commit> getCommitByIndex(int index) async {
     final response =
         await query(from: 'commits', where: fieldEquals('index', index));
-    return _commit(response.first.document);
+    return _commit(response.first);
   }
 
   Future<Commit> getLastCommit() async {
     final lastCommit =
         await query(from: 'commits', orderBy: orderBy('index', false));
-    return _commit(lastCommit.first.document);
+    return _commit(lastCommit.first);
   }
 
   Future<void> addCommit(String id, Map<String, dynamic> data) async {
@@ -169,7 +170,7 @@
           documentId: configuration);
       log('Configuration document $configuration -> $builder created');
     } else {
-      final originalBuilder = DataWrapper(record).getString('builder');
+      final originalBuilder = SafeDocument(record).getString('builder');
       if (originalBuilder != builder) {
         record.fields['builder'].stringValue = builder;
         await updateFields(record, ['builder']);
@@ -197,7 +198,7 @@
       documentsWritten++;
       return true;
     } else {
-      final data = DataWrapper(record);
+      final data = SafeDocument(record);
       final existingIndex = data.getInt('index');
       if (existingIndex != index) {
         throw ('Build $buildNumber of $builder had commit index '
@@ -249,19 +250,14 @@
         ]),
         limit: 5);
 
-    bool blamelistIncludesChange(RunQueryResponseElement response) {
-      final document = response.document;
-      if (document == null) return false;
-      final groupStart =
-          int.parse(document.fields['blamelist_start_index'].integerValue);
-      final groupEnd =
-          int.parse(document.fields['blamelist_end_index'].integerValue);
-      return startIndex <= groupEnd && endIndex >= groupStart;
+    bool blamelistIncludesChange(SafeDocument document) {
+      final before = endIndex < document.getInt('blamelist_start_index');
+      final after = startIndex > document.getInt('blamelist_end_index');
+      return !before && !after;
     }
 
     return snapshot
         .firstWhere(blamelistIncludesChange, orElse: () => null)
-        ?.document
         ?.name;
   }
 
@@ -280,7 +276,7 @@
     bool approved;
     await retryCommit(() async {
       final document = await getDocument(result);
-      final data = DataWrapper(document);
+      final data = SafeDocument(document);
       // Allow missing 'approved' field during transition period.
       approved = data.getBool('approved') ?? false;
       // Add the new configuration and narrow the blamelist.
@@ -341,15 +337,13 @@
   Future<List<Map<String, Value>>> findRevertedChanges(int index) async {
     final pinnedResults =
         await query(from: 'results', where: fieldEquals('pinned_index', index));
-    final results =
-        pinnedResults.map((response) => response.document.fields).toList();
+    final results = pinnedResults.map((response) => response.fields).toList();
     final unpinnedResults = await query(
         from: 'results', where: fieldEquals('blamelist_end_index', index));
-    for (final unpinnedResult in unpinnedResults) {
-      final data = DataWrapper(unpinnedResult.document);
+    for (final data in unpinnedResults) {
       if (data.getInt('blamelist_start_index') == index &&
           data.isNull('pinned_index')) {
-        results.add(unpinnedResult.document.fields);
+        results.add(data.fields);
       }
     }
     return results;
@@ -395,8 +389,8 @@
           ]),
           orderBy: orderBy('patchset', false),
           limit: 1);
-      final approved = previous.isNotEmpty &&
-          DataWrapper(previous.first.document).getBool('approved') == true;
+      final approved =
+          previous.isNotEmpty && previous.first.getBool('approved') == true;
 
       final document = Document()
         ..fields = taggedMap({
@@ -423,7 +417,7 @@
       documentsWritten++;
       return approved;
     } else {
-      final document = responses.first.document;
+      final document = responses.first;
       // Update the TryResult for this test, adding this configuration.
       final values = ArrayValue()..values = [taggedValue(configuration)];
       final addConfiguration = FieldTransform()
@@ -431,10 +425,10 @@
         ..appendMissingElements = values;
       await _executeWrite([
         Write()
-          ..update = document
+          ..update = document.toDocument()
           ..updateTransforms = [addConfiguration]
       ]);
-      return DataWrapper(document).getBool('approved') == true;
+      return document.getBool('approved') == true;
     }
   }
 
@@ -450,18 +444,17 @@
   /// Removes [configuration] from the active configurations and marks the
   /// active result inactive when we remove the last active config.
   Future<void> removeActiveConfiguration(
-      Document activeResult, String configuration) async {
-    final data = DataWrapper(activeResult);
-    final configurations = data.getList('active_configurations');
+      SafeDocument activeResult, String configuration) async {
+    final configurations = activeResult.getList('active_configurations');
     assert(configurations.contains(configuration));
     await removeArrayEntry(
         activeResult, 'active_configurations', taggedValue(configuration));
-    activeResult = await getDocument(activeResult.name);
-    if (DataWrapper(activeResult).getList('active_configurations').isEmpty) {
+    activeResult = SafeDocument(await getDocument(activeResult.name));
+    if (activeResult.getList('active_configurations').isEmpty) {
       activeResult.fields.remove('active_configurations');
       activeResult.fields.remove('active');
       final write = Write()
-        ..update = activeResult
+        ..update = activeResult.toDocument()
         ..updateMask =
             (DocumentMask()..fieldPaths = ['active', 'active_configurations']);
       await _executeWrite([write]);
@@ -469,7 +462,7 @@
   }
 
   Future<void> removeArrayEntry(
-      Document document, String fieldName, Value entry) async {
+      SafeDocument document, String fieldName, Value entry) async {
     await _executeWrite([
       Write()
         ..transform = (DocumentTransform()
@@ -482,7 +475,7 @@
     ]);
   }
 
-  Future<List<Document>> findActiveResults(
+  Future<List<SafeDocument>> findActiveResults(
       String name, String configuration) async {
     final results = await query(
         from: 'results',
@@ -497,9 +490,7 @@
         ...results
       ].join('\n'));
     }
-    return results
-        .map((RunQueryResponseElement result) => result.document)
-        .toList();
+    return results;
   }
 
   Future<bool> hasReview(String review) async {
@@ -578,19 +569,18 @@
         await query(from: 'comments', where: fieldEquals('review', review));
     if (comments.isEmpty) return;
     final writes = <Write>[];
-    for (final comment in comments) {
-      final document = comment.document;
+    for (final document in comments) {
       document.fields['blamelist_start_index'] = taggedValue(index);
       document.fields['blamelist_end_index'] = taggedValue(index);
       writes.add(Write()
-        ..update = document
+        ..update = document.toDocument()
         ..updateMask = (DocumentMask()
           ..fieldPaths = ['blamelist_start_index', 'blamelist_end_index']));
     }
     await _executeWrite(writes);
   }
 
-  Future<List<Map<String, Value>>> tryApprovals(int review) async {
+  Future<List<SafeDocument>> tryApprovals(int review) async {
     final patchsets = await query(
         from: 'patchsets',
         parent: 'reviews/$review',
@@ -599,20 +589,17 @@
     if (patchsets.isEmpty) {
       return [];
     }
-    final lastPatchsetGroup =
-        DataWrapper(patchsets.first.document).getInt('patchset_group');
-    final approvals = await query(
+    final lastPatchsetGroup = patchsets.first.getInt('patchset_group');
+    return query(
         from: 'try_results',
         where: compositeFilter([
           fieldEquals('approved', true),
           fieldEquals('review', review),
           fieldGreaterThanOrEqual('patchset', lastPatchsetGroup)
         ]));
-
-    return approvals.map((response) => response.document.fields).toList();
   }
 
-  Future<List<Map<String, Value>>> tryResults(
+  Future<List<SafeDocument>> tryResults(
       int review, String configuration) async {
     final patchsets = await query(
         from: 'patchsets',
@@ -622,18 +609,19 @@
     if (patchsets.isEmpty) {
       return [];
     }
-    final lastPatchsetGroup =
-        DataWrapper(patchsets.first.document).getInt('patchset_group');
-    final approvals = await query(
+    final lastPatchsetGroup = patchsets.first.getInt('patchset_group');
+    return query(
         from: 'try_results',
         where: compositeFilter([
           fieldEquals('review', review),
           arrayContains('configurations', configuration),
           fieldGreaterThanOrEqual('patchset', lastPatchsetGroup)
         ]));
-    return approvals.map((r) => r.document.fields).toList();
   }
 
+  // TODO(b/210111991): Remove Builder and TryBuilder records completely.
+  // They were only used to communicate success back to the builder
+  // when the processing was not on the builder.
   Future<void> completeBuilderRecord(
       String builder, int index, bool success) async {
     final path = '$documents/builds/$builder:$index';
@@ -643,24 +631,13 @@
 
   Future<void> _completeBuilderRecord(Document document, bool success) async {
     await retryCommit(() async {
-      final data = DataWrapper(document);
-      // TODO: Legacy support, remove if not needed anymore.
-      document.fields['processed_chunks'] = taggedValue(1);
-      document.fields['num_chunks'] = taggedValue(1);
-      document.fields['success'] =
-          taggedValue((data.getBool('success') ?? true) && success);
-      document.fields['completed'] = taggedValue(true);
-
+      document.fields['success'] = taggedValue(success);
       final write = Write()
         ..update = document
         ..updateMask = (DocumentMask()
           ..fieldPaths = [
-            'processed_chunks',
-            'num_chunks',
             'success',
-            'completed',
-          ])
-        ..currentDocument = (Precondition()..updateTime = document.updateTime);
+          ]);
       return write;
     });
   }
diff --git a/builder/lib/src/firestore_helpers.dart b/builder/lib/src/firestore_helpers.dart
index c4c2f99..b4bea17 100644
--- a/builder/lib/src/firestore_helpers.dart
+++ b/builder/lib/src/firestore_helpers.dart
@@ -4,6 +4,43 @@
 
 import 'package:googleapis/firestore/v1.dart';
 
+class SafeDocument {
+  final String /*!*/ name;
+  final Map<String, Value> /*!*/ fields;
+
+  SafeDocument(Document document)
+      : name = document.name /*!*/,
+        fields = document.fields /*!*/;
+
+  Document toDocument() => Document(name: name, fields: fields);
+  int getInt(String name) {
+    final value = fields[name]?.integerValue;
+    if (value == null) {
+      return null;
+    }
+    return int.parse(value);
+  }
+
+  String getString(String name) {
+    return fields[name]?.stringValue;
+  }
+
+  bool getBool(String name) {
+    return fields[name]?.booleanValue;
+  }
+
+  List<dynamic> getList(String name) {
+    final arrayValue = fields[name]?.arrayValue;
+    if (arrayValue == null) return null;
+    return arrayValue.values?.map(getValue)?.toList() ?? [];
+  }
+
+  bool isNull(String name) {
+    return !fields.containsKey(name) ||
+        fields['name'].nullValue == 'NULL_VALUE';
+  }
+}
+
 Map<String, Value> taggedMap(Map<String, dynamic> fields) {
   return fields.map((key, value) => MapEntry(key, taggedValue(value)));
 }
@@ -106,30 +143,4 @@
   final Map<String, Value> fields;
   DataWrapper(Document document) : fields = document.fields;
   DataWrapper.fields(this.fields);
-  int getInt(String name) {
-    final value = fields[name]?.integerValue;
-    if (value == null) {
-      return null;
-    }
-    return int.parse(value);
-  }
-
-  String getString(String name) {
-    return fields[name]?.stringValue;
-  }
-
-  bool getBool(String name) {
-    return fields[name]?.booleanValue;
-  }
-
-  List<dynamic> getList(String name) {
-    final arrayValue = fields[name]?.arrayValue;
-    if (arrayValue == null) return null;
-    return arrayValue.values?.map(getValue)?.toList() ?? [];
-  }
-
-  bool isNull(String name) {
-    return !fields.containsKey(name) ||
-        fields['name'].nullValue == 'NULL_VALUE';
-  }
 }
diff --git a/builder/lib/src/tryjob.dart b/builder/lib/src/tryjob.dart
index ae7ca32..c07ef13 100644
--- a/builder/lib/src/tryjob.dart
+++ b/builder/lib/src/tryjob.dart
@@ -3,7 +3,6 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import 'package:collection/collection.dart';
-import 'package:googleapis/firestore/v1.dart';
 import 'package:http/http.dart' as http show BaseClient;
 import 'package:pool/pool.dart';
 
@@ -66,8 +65,8 @@
   final TestNameLock testNameLock = TestNameLock();
   String baseRevision;
   bool success = true;
-  List<Map<String, Value>> landedResults;
-  Map<String, Map<String, Value>> lastLandedResultByName = {};
+  List<SafeDocument> landedResults;
+  Map<String, SafeDocument> lastLandedResultByName = {};
   final String buildbucketID;
 
   Tryjob(this.info, this.buildbucketID, this.baseRevision, this.commits,
@@ -83,7 +82,7 @@
   bool isNotLandedResult(Map<String, dynamic> change) {
     return !lastLandedResultByName.containsKey(change[fName]) ||
         change[fResult] !=
-            lastLandedResultByName[change[fName]][fResult].stringValue;
+            lastLandedResultByName[change[fName]].getString(fResult);
   }
 
   Future<void> process(List<Map<String, dynamic>> results) async {
@@ -97,8 +96,7 @@
         landedResults = await fetchLandedResults(configuration);
         // Map will contain the last result with each name.
         lastLandedResultByName = {
-          for (final result in landedResults)
-            getValue(result[fName]) as String: result
+          for (final result in landedResults) result.getString(fName): result
         };
       }
       final changes =
@@ -137,15 +135,13 @@
     }
   }
 
-  Future<List<Map<String, Value>>> fetchLandedResults(
-      String configuration) async {
+  Future<List<SafeDocument>> fetchLandedResults(String configuration) async {
     final resultsBase = await commits.getCommit(info.previousCommitHash);
     final rebaseBase = await commits.getCommit(baseRevision);
-    final results = <Map<String, Value>>[];
     if (resultsBase.index > rebaseBase.index) {
       print('Try build is rebased on $baseRevision, which is before '
           'the commit ${info.previousCommitHash} with CI comparison results');
-      return results;
+      return [];
     }
     final reviews = <int>[];
     for (var index = resultsBase.index + 1;
@@ -156,9 +152,9 @@
         reviews.add(commit.review);
       }
     }
-    for (final landedReview in reviews) {
-      results.addAll(await firestore.tryResults(landedReview, configuration));
-    }
-    return results;
+    return [
+      for (final landedReview in reviews)
+        ...await firestore.tryResults(landedReview, configuration)
+    ];
   }
 }
diff --git a/builder/test/approvals_test.dart b/builder/test/approvals_test.dart
index 8960fc3..7cc792b 100644
--- a/builder/test/approvals_test.dart
+++ b/builder/test/approvals_test.dart
@@ -59,8 +59,8 @@
 }
 
 Future<void> removeBuildersAndResults() async {
-  Future<void> deleteDocuments(RunQueryResponse response) async {
-    for (final document in response.map((r) => r.document)) {
+  Future<void> deleteDocuments(List<SafeDocument> documents) async {
+    for (final document in documents) {
       await firestore.deleteDocument(document.name);
     }
   }
@@ -94,10 +94,10 @@
       orderBy: orderBy('landed_index', false),
       where: fieldLessThanOrEqual('landed_index', startIndex),
       limit: 2);
-  final firstReview = reviews.first.document;
+  final firstReview = reviews.first;
   index1 = firstReview.fields['landed_index'].integerValue;
   review = firstReview.name.split('/').last;
-  final secondReview = reviews.last.document;
+  final secondReview = reviews.last;
   index2 = secondReview.fields['landed_index'].integerValue;
   review2 = secondReview.name.split('/').last;
   index3 = (int.parse(index2) - 1).toString();
@@ -108,7 +108,7 @@
     parent: 'reviews/$review',
     orderBy: orderBy('number', true),
   );
-  final patchsetFields = patchsets.last.document.fields;
+  final patchsetFields = patchsets.last.fields;
   lastPatchset = patchsetFields['number'].integerValue;
   lastPatchsetRef = 'refs/changes/$review/$lastPatchset';
   patchsetGroup = patchsetFields['patchset_group'].integerValue;
@@ -120,7 +120,7 @@
     parent: 'reviews/$review2',
     orderBy: orderBy('number', true),
   );
-  patchset2 = patchsets2.last.document.fields['number'].integerValue;
+  patchset2 = patchsets2.last.fields['number'].integerValue;
   patchset2Ref = 'refs/changes/$review/$patchset2';
 
   // Get commit hashes for the landed reviews, and for a commit before them
@@ -131,7 +131,6 @@
               where: fieldEquals('index', int.parse(index)),
               limit: 1))
           .first
-          .document
           .name
           .split('/')
           .last
@@ -234,13 +233,13 @@
     await makeTryjob('approvals', change1).process([change1]);
     var documents = await firestore.query(
         from: 'try_results', where: fieldEquals('name', 'approvals_test'));
-    await firestore.approveResult(documents.single.document);
+    await firestore.approveResult(documents.single.toDocument());
     final change2 = makeTryChange('approvals', newFailure, patchsetGroupRef,
         testName: 'approvals_2');
     await makeTryjob('approvals', change2).process([change2]);
     documents = await firestore.query(
         from: 'try_results', where: fieldEquals('name', 'approvals_2_test'));
-    await firestore.approveResult(documents.single.document);
+    await firestore.approveResult(documents.single.toDocument());
 
     final change3 = makeChange('approvals', newFailure, commit1, commit4);
     final change3a = makeChange('approvals', newFailure, commit1, commit4)
@@ -280,16 +279,16 @@
         from: 'comments',
         where: fieldEquals('review', int.parse(reviewWithComments)));
     final landedIndex =
-        commentsQuery.first.document.fields[fBlamelistStartIndex].integerValue;
+        commentsQuery.first.fields[fBlamelistStartIndex].integerValue;
     for (final item in commentsQuery) {
-      final fields = item.document.fields;
+      final fields = item.fields;
       expect(fields[fBlamelistStartIndex].integerValue, landedIndex);
       expect(fields[fBlamelistEndIndex].integerValue, landedIndex);
       expect(fields[fReview].integerValue, reviewWithComments);
       fields.remove(fBlamelistStartIndex);
       fields.remove(fBlamelistEndIndex);
       await firestore.updateFields(
-          item.document, [fBlamelistStartIndex, fBlamelistEndIndex]);
+          item.toDocument(), [fBlamelistStartIndex, fBlamelistEndIndex]);
     }
     var reviewDocument = await firestore
         .getDocument('${firestore.documents}/reviews/$reviewWithComments');
@@ -305,7 +304,7 @@
         from: 'comments',
         where: fieldEquals('review', int.parse(reviewWithComments)));
     for (final item in commentsQuery) {
-      final fields = item.document.fields;
+      final fields = item.fields;
       expect(fields[fBlamelistStartIndex].integerValue, landedIndex);
       expect(fields[fBlamelistEndIndex].integerValue, landedIndex);
       expect(fields[fReview].integerValue, reviewWithComments);
@@ -321,12 +320,12 @@
   final buildDocuments = await firestore.query(
       from: 'try_builds', where: fieldEquals('buildbucket_id', buildbucketId));
   expect(buildDocuments.length, 1);
-  final data = untagMap(buildDocuments.single.document.fields);
-  expect(data['success'], success);
+  final document = buildDocuments.single;
+  expect(document.getBool('success'), success);
   if (truncated != null) {
-    expect(data['truncated'], truncated);
+    expect(document.getBool('truncated'), truncated);
   } else {
-    expect(data.containsKey('truncated'), isFalse);
+    expect(document.fields.containsKey('truncated'), isFalse);
   }
 }
 
diff --git a/builder/test/fakes.dart b/builder/test/fakes.dart
index 4d603c2..8502ea7 100644
--- a/builder/test/fakes.dart
+++ b/builder/test/fakes.dart
@@ -102,16 +102,16 @@
   }
 
   @override
-  Future<List<Document>> findActiveResults(
+  Future<List<SafeDocument>> findActiveResults(
       String name, String configuration) async {
     return [
       for (final id in results.keys)
         if (results[id][fName] == name &&
             results[id][fActiveConfigurations] != null &&
             results[id][fActiveConfigurations].contains(configuration))
-          Document()
+          SafeDocument(Document()
             ..fields = taggedMap(Map.from(results[id]))
-            ..name = id
+            ..name = id)
     ];
   }
 
@@ -153,7 +153,7 @@
 
   @override
   Future<void> removeActiveConfiguration(
-      Document activeResult, String configuration) async {
+      SafeDocument activeResult, String configuration) async {
     final result = Map<String, dynamic>.from(results[activeResult.name]);
     result[fActiveConfigurations] = List.from(result[fActiveConfigurations])
       ..remove(configuration);
@@ -176,22 +176,28 @@
   }
 
   @override
-  Future<List<Map<String, Value>>> tryApprovals(int review) async {
+  Future<List<SafeDocument>> tryApprovals(int review) async {
     return fakeTryResults
         .where(
             (result) => result[fReview] == review && result[fApproved] == true)
         .map(taggedMap)
+        .map((fields) => SafeDocument(Document()
+          ..fields = fields
+          ..name = ''))
         .toList();
   }
 
   @override
-  Future<List<Map<String, Value>>> tryResults(
+  Future<List<SafeDocument>> tryResults(
       int review, String configuration) async {
     return fakeTryResults
         .where((result) =>
             result[fReview] == review &&
             result[fConfigurations].contains(configuration))
         .map(taggedMap)
+        .map((fields) => SafeDocument(Document()
+          ..fields = fields
+          ..name = ''))
         .toList();
   }
 
diff --git a/builder/test/firestore_test.dart b/builder/test/firestore_test.dart
index beb50d8..57f0516 100644
--- a/builder/test/firestore_test.dart
+++ b/builder/test/firestore_test.dart
@@ -53,24 +53,19 @@
       var snapshot = await firestore.query(
           from: 'try_builds', where: fieldEquals('review', testReview));
       for (final doc in snapshot) {
-        if (doc.document != null) {
-          await firestore.deleteDocument(doc.document.name);
-        }
+        await firestore.deleteDocument(doc.name);
       }
+
       snapshot = await firestore.query(
           from: 'patchsets', parent: 'reviews/$testReview/');
       for (final doc in snapshot) {
-        if (doc.document != null) {
-          await firestore.deleteDocument(doc.document.name);
-        }
+        await firestore.deleteDocument(doc.name);
       }
       snapshot = await firestore.query(
           from: 'results',
           where: fieldEquals('name', removeActiveConfigurationTestName));
       for (final doc in snapshot) {
-        if (doc.document != null) {
-          await firestore.deleteDocument(doc.document.name);
-        }
+        await firestore.deleteDocument(doc.name);
       }
       await firestore.deleteDocument(testReviewDocument);
     });
@@ -161,7 +156,7 @@
             fieldLessThanOrEqual('patchset', 2)
           ]));
       for (final response in snapshot) {
-        await firestore.approveResult(response.document);
+        await firestore.approveResult(response.toDocument());
         //await firestore.updateDocument(response.document.name, {'approved': taggedValue(true)});
       }
 
@@ -172,7 +167,7 @@
       tryResult2['approved'] = true;
       tryResult2.remove('configuration');
       expect(1, approvals.length);
-      final approval = untagMap(approvals.single);
+      final approval = untagMap(approvals.single.fields);
       expect(approval, tryResult2);
     });
   });
diff --git a/builder/test/tryjob_test.dart b/builder/test/tryjob_test.dart
index c1bbb3e..08718ff 100644
--- a/builder/test/tryjob_test.dart
+++ b/builder/test/tryjob_test.dart
@@ -37,8 +37,8 @@
 }
 
 Future<void> removeTryBuildersAndResults() async {
-  Future<void> deleteDocuments(RunQueryResponse response) async {
-    for (final document in response.map((r) => r.document)) {
+  Future<void> deleteDocuments(List<SafeDocument> documents) async {
+    for (final document in documents) {
       await firestore.deleteDocument(document.name);
     }
   }
@@ -60,10 +60,10 @@
       orderBy: orderBy('landed_index', false),
       where: fieldLessThanOrEqual('landed_index', startIndex),
       limit: 2);
-  final firstReview = reviews.first.document;
+  final firstReview = reviews.first;
   final String index = firstReview.fields['landed_index'].integerValue;
   final String review = firstReview.name.split('/').last;
-  final secondReview = reviews.last.document;
+  final secondReview = reviews.last;
   final String landedIndex = secondReview.fields['landed_index'].integerValue;
   final String landedReview = secondReview.name.split('/').last;
   // expect(int.parse(index), greaterThan(int.parse(landedIndex)));
@@ -74,15 +74,14 @@
     parent: 'reviews/$review',
     orderBy: orderBy('number', true),
   );
-  final patchset = patchsets.last.document.fields['number'].integerValue;
+  final patchset = patchsets.last.fields['number'].integerValue;
   final previousPatchset = '1';
   final landedPatchsets = await firestore.query(
     from: 'patchsets',
     parent: 'reviews/$landedReview',
     orderBy: orderBy('number', true),
   );
-  final landedPatchset =
-      landedPatchsets.last.document.fields['number'].integerValue;
+  final landedPatchset = landedPatchsets.last.fields['number'].integerValue;
 
   // Get commit hashes for the landed reviews, and for a commit before them
   var commits = {
@@ -92,7 +91,6 @@
               where: fieldEquals('index', int.parse(index)),
               limit: 1))
           .first
-          .document
           .name
           .split('/')
           .last
@@ -167,14 +165,11 @@
   final buildDocuments = await firestore.query(
       from: 'try_builds', where: fieldEquals('buildbucket_id', buildbucketId));
   expect(buildDocuments.length, 1);
-  expect(
-      buildDocuments.single.document.fields['success'].booleanValue, success);
+  expect(buildDocuments.single.fields['success'].booleanValue, success);
   if (truncated != null) {
-    expect(buildDocuments.single.document.fields['truncated'].booleanValue,
-        truncated);
+    expect(buildDocuments.single.fields['truncated'].booleanValue, truncated);
   } else {
-    expect(buildDocuments.single.document.fields.containsKey('truncated'),
-        isFalse);
+    expect(buildDocuments.single.fields.containsKey('truncated'), isFalse);
   }
 }
 
@@ -215,10 +210,7 @@
     final result = await firestore.query(
         from: 'try_results', where: fieldEquals('name', 'failure_test'));
     expect(result.length, 1);
-    expect(
-        result
-            .single.document.fields['configurations'].arrayValue.values.length,
-        2);
+    expect(result.single.getList('configurations').length, 2);
   });
 
   test('landedFailure', () async {