Refactor Tryjobs & Reviews to use extension types (Part 2)

Introduce TryResultRecord, TryBuildRecord, ReviewRecord, PatchsetRecord, and CommentRecord extension types. Refactor tryjob.dart and tryjob/review Firestore service methods to use these new extension types. Refactor approvals_test.dart, tryjob_test.dart, and firestore_test.dart to test with the new extension types. Keep commits using Commit class for now, to be refactored in Part 3. Rename runQuery to query in FirestoreService and update all tests to use it. Add meta dependency for @visibleForTesting.

TAG=agy

CONV=a2878fa6-b79f-4b1d-9635-654148d5ceb2
diff --git a/builder/lib/src/builder.dart b/builder/lib/src/builder.dart
index f41fa45..c913196 100644
--- a/builder/lib/src/builder.dart
+++ b/builder/lib/src/builder.dart
@@ -136,7 +136,7 @@
       if (review != null) {
         tryApprovals.addAll({
           for (final result in await firestore.tryApprovals(review))
-            testResult(result.fields): index,
+            result.testResult: index,
         });
       }
       if (reverted != null) {
diff --git a/builder/lib/src/data.dart b/builder/lib/src/data.dart
index b1baaa9..fd2068c 100644
--- a/builder/lib/src/data.dart
+++ b/builder/lib/src/data.dart
@@ -17,10 +17,15 @@
   String get expected => doc.fields![fExpected]!.stringValue!;
   int get blamelistStartIndex =>
       int.parse(doc.fields![fBlamelistStartIndex]!.integerValue!);
+  set blamelistStartIndex(int value) =>
+      doc.fields![fBlamelistStartIndex] = taggedValue(value);
   int get blamelistEndIndex =>
       int.parse(doc.fields![fBlamelistEndIndex]!.integerValue!);
+  set blamelistEndIndex(int value) =>
+      doc.fields![fBlamelistEndIndex] = taggedValue(value);
   bool get approved => doc.fields![fApproved]?.booleanValue ?? false;
   bool get active => doc.fields![fActive]?.booleanValue ?? false;
+  set active(bool value) => doc.fields![fActive] = taggedValue(value);
   List<String> get configurations => doc
       .fields![fConfigurations]!
       .arrayValue!
@@ -86,6 +91,81 @@
   }
 }
 
+extension type TryResultRecord(Document doc) {
+  String get name => doc.fields![fName]!.stringValue!;
+  String get result => doc.fields![fResult]!.stringValue!;
+  String get previousResult => doc.fields![fPreviousResult]!.stringValue!;
+  String get expected => doc.fields![fExpected]!.stringValue!;
+  int get review => int.parse(doc.fields![fReview]!.integerValue!);
+  int get patchset => int.parse(doc.fields!['patchset']!.integerValue!);
+  bool get approved => doc.fields![fApproved]?.booleanValue ?? false;
+  List<String> get configurations => doc
+      .fields![fConfigurations]!
+      .arrayValue!
+      .values!
+      .map((v) => v.stringValue!)
+      .toList();
+
+  String get testResult => [name, result, previousResult, expected].join(' ');
+}
+
+extension type TryBuildRecord(Document doc) {
+  String get builder => doc.fields!['builder']!.stringValue!;
+  int get buildNumber => int.parse(doc.fields!['build_number']!.integerValue!);
+  String get buildbucketId => doc.fields!['buildbucket_id']!.stringValue!;
+  int get review => int.parse(doc.fields![fReview]!.integerValue!);
+  int get patchset => int.parse(doc.fields!['patchset']!.integerValue!);
+  bool get success => doc.fields!['success']?.booleanValue ?? false;
+  bool get completed => doc.fields!['completed']?.booleanValue ?? false;
+  bool get truncated => doc.fields!['truncated']?.booleanValue ?? false;
+}
+
+extension type BuildRecord(Document doc) {
+  String get builder => doc.fields!['builder']!.stringValue!;
+  int get buildNumber => int.parse(doc.fields!['build_number']!.integerValue!);
+  int get index => int.parse(doc.fields!['index']!.integerValue!);
+  bool get success => doc.fields!['success']?.booleanValue ?? false;
+  bool get completed => doc.fields!['completed']?.booleanValue ?? false;
+}
+
+extension type ReviewRecord(Document doc) {
+  String get review => doc.name!.split('/').last;
+  String get subject => doc.fields!['subject']!.stringValue!;
+  int? get landedIndex =>
+      int.tryParse(doc.fields!['landed_index']?.integerValue ?? '');
+  set landedIndex(int? value) =>
+      doc.fields!['landed_index'] = taggedValue(value);
+  String? get revertOf => doc.fields!['revert_of']?.stringValue;
+}
+
+extension type PatchsetRecord(Document doc) {
+  int get number => int.parse(doc.fields!['number']!.integerValue!);
+  int get patchsetGroup =>
+      int.parse(doc.fields!['patchset_group']!.integerValue!);
+  String get kind => doc.fields!['kind']!.stringValue!;
+  String? get description => doc.fields!['description']?.stringValue;
+}
+
+extension type CommentRecord(Document doc) {
+  String get id => doc.name!.split('/').last;
+  String get author => doc.fields!['author']!.stringValue!;
+  String get comment => doc.fields!['comment']!.stringValue!;
+  int get review => int.parse(doc.fields![fReview]!.integerValue!);
+  int? get blamelistStartIndex =>
+      int.tryParse(doc.fields![fBlamelistStartIndex]?.integerValue ?? '');
+  set blamelistStartIndex(int? value) {
+    doc.fields![fBlamelistStartIndex] = taggedValue(value);
+  }
+
+  int? get blamelistEndIndex =>
+      int.tryParse(doc.fields![fBlamelistEndIndex]?.integerValue ?? '');
+  set blamelistEndIndex(int? value) {
+    doc.fields![fBlamelistEndIndex] = taggedValue(value);
+  }
+
+  bool get approved => doc.fields![fApproved]?.booleanValue ?? false;
+}
+
 extension type ConfigurationRecord(Document doc) {
   String get builder => doc.fields!['builder']!.stringValue!;
 }
diff --git a/builder/lib/src/firestore.dart b/builder/lib/src/firestore.dart
index e3a8713..566134c 100644
--- a/builder/lib/src/firestore.dart
+++ b/builder/lib/src/firestore.dart
@@ -10,6 +10,7 @@
 import 'package:collection/collection.dart' show IterableExtension;
 import 'package:googleapis/firestore/v1.dart';
 import 'package:http/http.dart' as http;
+import 'package:meta/meta.dart';
 
 import 'firestore_helpers.dart';
 
@@ -52,22 +53,6 @@
     }
   }
 
-  Future<List<SafeDocument>> query({
-    required String from,
-    Filter? where,
-    Order? orderBy,
-    int? limit,
-    String? parent,
-  }) async {
-    final query = StructuredQuery()
-      ..from = inCollection(from)
-      ..where = where
-      ..orderBy = orderBy != null ? [orderBy] : null
-      ..limit = limit;
-    final results = await runQuery(query, parent: parent);
-    return results.map((d) => SafeDocument(d)).toList();
-  }
-
   Future<List<T>> _query<T>({
     required String from,
     Filter? where,
@@ -80,26 +65,26 @@
       ..where = where
       ..orderBy = orderBy != null ? [orderBy] : null
       ..limit = limit;
-    final results = await runQuery(query, parent: parent);
+    final results = await this.query(query, parent: parent);
     return results.cast<T>();
   }
 
-  Future<Document> getDocument(String path) async {
+  Future<T> getDocument<T>(String path) async {
     try {
       final document = await firestore.projects.databases.documents.get(path);
       documentsFetched++;
-      return document;
+      return document as T;
     } on DetailedApiRequestError {
       log("Failed to get document '$path'");
       rethrow;
     }
   }
 
-  Future<Document?> getDocumentOrNull(String path) async {
+  Future<T?> getDocumentOrNull<T>(String path) async {
     try {
       final document = await firestore.projects.databases.documents.get(path);
       documentsFetched++;
-      return document;
+      return document as T;
     } on DetailedApiRequestError catch (e) {
       if (e.status == 404) {
         return null;
@@ -117,10 +102,8 @@
   String get database => 'projects/$project/databases/(default)';
   String get documents => '$database/documents';
 
-  Future<List<Document>> runQuery(
-    StructuredQuery query, {
-    String? parent,
-  }) async {
+  @visibleForTesting
+  Future<List<Document>> query(StructuredQuery query, {String? parent}) async {
     final request = RunQueryRequest()..structuredQuery = query;
     final parentPath = parent == null ? documents : '$documents/$parent';
     final queryResponse = await firestore.projects.databases.documents.runQuery(
@@ -152,12 +135,14 @@
   }
 
   Future<Commit?> getCommit(String hash) async {
-    final document = await getDocumentOrNull('$documents/commits/$hash');
-    return document != null ? Commit(hash, SafeDocument(document)) : null;
+    final document = await getDocumentOrNull<SafeDocument>(
+      '$documents/commits/$hash',
+    );
+    return document != null ? Commit(hash, document) : null;
   }
 
   Future<Commit> getCommitByIndex(int index) async {
-    final response = await query(
+    final response = await _query<SafeDocument>(
       from: 'commits',
       where: fieldEquals('index', index),
     );
@@ -165,7 +150,7 @@
   }
 
   Future<Commit> getLastCommit() async {
-    final lastCommit = await query(
+    final lastCommit = await _query<SafeDocument>(
       from: 'commits',
       orderBy: orderBy('index', false),
     );
@@ -194,7 +179,7 @@
 
   Future<void> updateConfiguration(String configuration, String builder) async {
     documentsWritten++;
-    final record = await getDocumentOrNull(
+    final record = await getDocumentOrNull<ConfigurationRecord>(
       '$documents/configurations/$configuration',
     );
     if (record == null) {
@@ -207,10 +192,10 @@
       );
       log('Configuration document $configuration -> $builder created');
     } else {
-      final originalBuilder = SafeDocument(record).getString('builder');
+      final originalBuilder = record.builder;
       if (originalBuilder != builder) {
-        record.fields!['builder']!.stringValue = builder;
-        await updateFields(record, ['builder']);
+        record.doc.fields!['builder']!.stringValue = builder;
+        await updateFields(record.doc, ['builder']);
         log(
           'Configuration document changed: $configuration -> $builder '
           '(was $originalBuilder)',
@@ -228,7 +213,9 @@
     int buildNumber,
     int index,
   ) async {
-    final record = await getDocumentOrNull('$documents/builds/$builder:$index');
+    final record = await getDocumentOrNull<BuildRecord>(
+      '$documents/builds/$builder:$index',
+    );
     if (record == null) {
       final newRecord = Document()
         ..fields = taggedMap({
@@ -245,13 +232,12 @@
       documentsWritten++;
       return true;
     } else {
-      final data = SafeDocument(record);
-      final existingIndex = data.getInt('index');
+      final existingIndex = record.index;
       if (existingIndex != index) {
         throw ('Build $buildNumber of $builder had commit index '
             '$existingIndex, should be $index.');
       }
-      return data.getBool('completed') != true;
+      return record.completed != true;
     }
   }
 
@@ -348,8 +334,7 @@
   }) async {
     late bool approved;
     await retryCommit(() async {
-      final document = await getDocument(result);
-      final data = ResultRecord(document);
+      final data = await getDocument<ResultRecord>(result);
       approved = data.approved;
       // Add the new configuration and narrow the blamelist.
       final newStart = max(startIndex, data.blamelistStartIndex);
@@ -361,16 +346,16 @@
         'blamelist_end_index',
         if (failure) 'active',
       ];
-      document.fields!['blamelist_start_index'] = taggedValue(newStart);
-      document.fields!['blamelist_end_index'] = taggedValue(newEnd);
+      data.blamelistStartIndex = newStart;
+      data.blamelistEndIndex = newEnd;
       if (failure) {
-        document.fields!['active'] = taggedValue(true);
+        data.active = true;
       }
       final addConfiguration = ArrayValue()
         ..values = [taggedValue(configuration)];
       final write = Write()
-        ..currentDocument = (Precondition()..updateTime = document.updateTime)
-        ..update = document
+        ..currentDocument = (Precondition()..updateTime = data.doc.updateTime)
+        ..update = data.doc
         ..updateMask = (DocumentMask()..fieldPaths = updates)
         ..updateTransforms = [
           FieldTransform()
@@ -438,7 +423,7 @@
     final configuration = change.configuration;
 
     // Find an existing result record for this test on this patchset.
-    final responses = await query(
+    final responses = await _query<TryResultRecord>(
       from: 'try_results',
       where: compositeFilter([
         fieldEquals('review', review),
@@ -459,7 +444,7 @@
     //                  or there is no response at all.
     if (responses.isEmpty) {
       // Is the previous result for this test on this review approved?
-      final previous = await query(
+      final previous = await _query<TryResultRecord>(
         from: 'try_results',
         where: compositeFilter([
           fieldEquals('review', review),
@@ -471,8 +456,7 @@
         orderBy: orderBy('patchset', false),
         limit: 1,
       );
-      final approved =
-          previous.isNotEmpty && previous.first.getBool('approved') == true;
+      final approved = previous.isNotEmpty && previous.first.approved;
 
       final document = Document()
         ..fields = taggedMap({
@@ -503,7 +487,7 @@
       documentsWritten++;
       return approved;
     } else {
-      final document = responses.first;
+      final tryResult = responses.first;
       // Update the TryResult for this test, adding this configuration.
       final values = ArrayValue()..values = [taggedValue(configuration)];
       final addConfiguration = FieldTransform()
@@ -511,10 +495,10 @@
         ..appendMissingElements = values;
       await _executeWrite([
         Write()
-          ..update = document.toDocument()
+          ..update = tryResult.doc
           ..updateTransforms = [addConfiguration],
       ]);
-      return document.getBool('approved') == true;
+      return tryResult.approved;
     }
   }
 
@@ -540,8 +524,7 @@
       'active_configurations',
       taggedValue(configuration),
     );
-    final document = await getDocument(activeResult.doc.name!);
-    activeResult = ResultRecord(document);
+    activeResult = await getDocument<ResultRecord>(activeResult.doc.name!);
     if (activeResult.activeConfigurations?.isEmpty == true) {
       activeResult.doc.fields!.remove('active_configurations');
       activeResult.doc.fields!.remove('active');
@@ -628,7 +611,7 @@
   }
 
   Future<bool> documentExists(String name) async {
-    return (await getDocumentOrNull(name) != null);
+    return (await getDocumentOrNull<Document>(name) != null);
   }
 
   Future _executeWrite(List<Write> writes) async {
@@ -673,32 +656,36 @@
   /// or if there is no record for the review in the database.  Reviews with no
   /// test failures have no record, and don't need to be linked when landing.
   Future<bool> reviewIsLanded(int review) async {
-    final document = await getDocumentOrNull('$documents/reviews/$review');
+    final document = await getDocumentOrNull<ReviewRecord>(
+      '$documents/reviews/$review',
+    );
     if (document == null) {
       return true;
     }
-    return document.fields!.containsKey('landed_index');
+    return document.doc.fields!.containsKey('landed_index');
   }
 
   Future<void> linkReviewToCommit(int review, int index) async {
-    final document = await getDocument('$documents/reviews/$review');
-    document.fields!['landed_index'] = taggedValue(index);
-    await updateFields(document, ['landed_index']);
+    final document = await getDocument<ReviewRecord>(
+      '$documents/reviews/$review',
+    );
+    document.landedIndex = index;
+    await updateFields(document.doc, ['landed_index']);
   }
 
   Future<void> linkCommentsToCommit(int review, int index) async {
-    final comments = await query(
+    final comments = await _query<CommentRecord>(
       from: 'comments',
       where: fieldEquals('review', review),
     );
     if (comments.isEmpty) return;
     final writes = <Write>[];
-    for (final document in comments) {
-      document.fields['blamelist_start_index'] = taggedValue(index);
-      document.fields['blamelist_end_index'] = taggedValue(index);
+    for (final comment in comments) {
+      comment.blamelistStartIndex = index;
+      comment.blamelistEndIndex = index;
       writes.add(
         Write()
-          ..update = document.toDocument()
+          ..update = comment.doc
           ..updateMask = (DocumentMask()
             ..fieldPaths = ['blamelist_start_index', 'blamelist_end_index']),
       );
@@ -706,8 +693,8 @@
     await _executeWrite(writes);
   }
 
-  Future<List<SafeDocument>> tryApprovals(int review) async {
-    final patchsets = await query(
+  Future<List<TryResultRecord>> tryApprovals(int review) async {
+    final patchsets = await _query<PatchsetRecord>(
       from: 'patchsets',
       parent: 'reviews/$review',
       orderBy: orderBy('number', false),
@@ -716,8 +703,8 @@
     if (patchsets.isEmpty) {
       return [];
     }
-    final lastPatchsetGroup = patchsets.first.getInt('patchset_group');
-    return query(
+    final lastPatchsetGroup = patchsets.first.patchsetGroup;
+    return _query<TryResultRecord>(
       from: 'try_results',
       where: compositeFilter([
         fieldEquals('approved', true),
@@ -727,11 +714,11 @@
     );
   }
 
-  Future<List<SafeDocument>> tryResults(
+  Future<List<TryResultRecord>> tryResults(
     int review,
     String configuration,
   ) async {
-    final patchsets = await query(
+    final patchsets = await _query<PatchsetRecord>(
       from: 'patchsets',
       parent: 'reviews/$review',
       orderBy: orderBy('number', false),
@@ -740,8 +727,8 @@
     if (patchsets.isEmpty) {
       return [];
     }
-    final lastPatchsetGroup = patchsets.first.getInt('patchset_group');
-    return query(
+    final lastPatchsetGroup = patchsets.first.patchsetGroup;
+    return _query<TryResultRecord>(
       from: 'try_results',
       where: compositeFilter([
         fieldEquals('review', review),
@@ -760,7 +747,7 @@
     bool success,
   ) async {
     final path = '$documents/builds/$builder:$index';
-    final document = await getDocument(path);
+    final document = await getDocument<Document>(path);
     await _completeBuilderRecord(document, success);
   }
 
diff --git a/builder/lib/src/firestore_helpers.dart b/builder/lib/src/firestore_helpers.dart
index 9e1903f..1f5da08 100644
--- a/builder/lib/src/firestore_helpers.dart
+++ b/builder/lib/src/firestore_helpers.dart
@@ -6,15 +6,12 @@
 
 export 'data.dart';
 
-class SafeDocument {
-  final String name;
-  final Map<String, Value> fields;
+extension type SafeDocument(Document doc) {
+  String get name => doc.name!;
+  Map<String, Value> get fields => doc.fields!;
 
-  SafeDocument(Document document)
-    : name = document.name!,
-      fields = document.fields!;
+  Document toDocument() => doc;
 
-  Document toDocument() => Document(name: name, fields: fields);
   int? getInt(String name) {
     final value = fields[name]?.integerValue;
     if (value == null) {
diff --git a/builder/lib/src/tryjob.dart b/builder/lib/src/tryjob.dart
index 15fefc9..17ccc08 100644
--- a/builder/lib/src/tryjob.dart
+++ b/builder/lib/src/tryjob.dart
@@ -66,8 +66,8 @@
   final TestNameLock testNameLock = TestNameLock();
   String baseRevision;
   bool success = true;
-  late List<SafeDocument> landedResults;
-  Map<String, SafeDocument> lastLandedResultByName = {};
+  late List<TryResultRecord> landedResults;
+  Map<String, TryResultRecord> lastLandedResultByName = {};
   final String buildbucketID;
 
   Tryjob(
@@ -91,8 +91,7 @@
   }
 
   bool isNotLandedResult(ChangeRecord change) {
-    return change.result !=
-        lastLandedResultByName[change.name]?.getString(fResult);
+    return change.result != lastLandedResultByName[change.name]?.result;
   }
 
   Future<BuildStatus> process(List<ChangeRecord> results) async {
@@ -108,7 +107,7 @@
         landedResults = await fetchLandedResults(configuration);
         // Map will contain the last result with each name.
         lastLandedResultByName = {
-          for (final result in landedResults) result.getString(fName): result,
+          for (final result in landedResults) result.name: result,
         };
       }
       final changes = resultsByConfiguration[configuration]!.where(
@@ -159,7 +158,7 @@
     }
   }
 
-  Future<List<SafeDocument>> fetchLandedResults(String configuration) async {
+  Future<List<TryResultRecord>> fetchLandedResults(String configuration) async {
     final resultsBase = await commits.getCommit(info.previousCommitHash!);
     final rebaseBase = await commits.getCommit(baseRevision);
     if (resultsBase.index > rebaseBase.index) {
diff --git a/builder/pubspec.lock b/builder/pubspec.lock
index d20320f..ce1d2ea 100644
--- a/builder/pubspec.lock
+++ b/builder/pubspec.lock
@@ -202,13 +202,13 @@
     source: hosted
     version: "0.12.17"
   meta:
-    dependency: transitive
+    dependency: "direct main"
     description:
       name: meta
-      sha256: "23f08335362185a5ea2ad3a4e597f1375e78bce8a040df5c600c8d3552ef2394"
+      sha256: c82594181e3312f3d0695fc95aaaf7758d75b8d4ae2bbecf223b9fd5109a059d
       url: "https://pub.dev"
     source: hosted
-    version: "1.17.0"
+    version: "1.18.3"
   mime:
     dependency: transitive
     description:
diff --git a/builder/pubspec.yaml b/builder/pubspec.yaml
index 98b6b8e..828b175 100644
--- a/builder/pubspec.yaml
+++ b/builder/pubspec.yaml
@@ -15,6 +15,7 @@
   retry: ^3.1.0
   collection: ^1.15.0
   glob: ^2.1.0
+  meta: ^1.18.3
 
 dev_dependencies:
   lints: ^6.0.0
diff --git a/builder/test/approvals_test.dart b/builder/test/approvals_test.dart
index 2c6b518..9c9eb42 100644
--- a/builder/test/approvals_test.dart
+++ b/builder/test/approvals_test.dart
@@ -32,20 +32,20 @@
 // These globals are populated by loadTestCommits().
 const testCommitsStart = 80801;
 const reviewWithComments = '215021';
-late final String index1; // Index of the final commit in the test range
+late final int index1; // Index of the final commit in the test range
 late final String commit1; // Hash of that commit
 late final String review; // CL number of that commit's Gerrit review
-late final String lastPatchset; // Final patchset in that review
+late final int lastPatchset; // Final patchset in that review
 late final String lastPatchsetRef; // 'refs/changes/[review]/[patchset]'
-late final String patchsetGroup; // First patchset in the final patchset group
+late final int patchsetGroup; // First patchset in the final patchset group
 late final String patchsetGroupRef;
-late final String earlyPatchset; // Patchset not in the final patchset group
+late final int earlyPatchset; // Patchset not in the final patchset group
 late final String earlyPatchsetRef;
 // Earlier commit with a review
-late final String index2;
+late final int index2;
 late final String commit2;
 late final String review2;
-late final String patchset2;
+late final int patchset2;
 late final String patchset2Ref;
 // Commits before commit2
 late final String index3;
@@ -62,46 +62,54 @@
 }
 
 Future<void> removeBuildersAndResults() async {
-  Future<void> deleteDocuments(List<SafeDocument> documents) async {
+  Future<void> deleteDocuments(List<Document> documents) async {
     for (final document in documents) {
-      await firestore.deleteDocument(document.name);
+      await firestore.deleteDocument(document.name!);
     }
   }
 
   for (final test in testsToRemove) {
     await deleteDocuments(
       await firestore.query(
-        from: 'try_results',
-        where: fieldEquals('name', test),
+        StructuredQuery()
+          ..from = inCollection('try_results')
+          ..where = fieldEquals('name', test),
       ),
     );
   }
   for (final test in testsToRemove) {
     await deleteDocuments(
-      await firestore.query(from: 'results', where: fieldEquals('name', test)),
-    );
-  }
-  for (final builder in buildersToRemove) {
-    await deleteDocuments(
       await firestore.query(
-        from: 'try_builds',
-        where: fieldEquals('builder', builder),
+        StructuredQuery()
+          ..from = inCollection('results')
+          ..where = fieldEquals('name', test),
       ),
     );
   }
   for (final builder in buildersToRemove) {
     await deleteDocuments(
       await firestore.query(
-        from: 'builds',
-        where: fieldEquals('builder', builder),
+        StructuredQuery()
+          ..from = inCollection('try_builds')
+          ..where = fieldEquals('builder', builder),
       ),
     );
   }
   for (final builder in buildersToRemove) {
     await deleteDocuments(
       await firestore.query(
-        from: 'configurations',
-        where: fieldEquals('builder', builder),
+        StructuredQuery()
+          ..from = inCollection('builds')
+          ..where = fieldEquals('builder', builder),
+      ),
+    );
+  }
+  for (final builder in buildersToRemove) {
+    await deleteDocuments(
+      await firestore.query(
+        StructuredQuery()
+          ..from = inCollection('configurations')
+          ..where = fieldEquals('builder', builder),
       ),
     );
   }
@@ -110,51 +118,55 @@
 Future<void> loadTestCommits(int startIndex) async {
   // Get review data for the last two landed CLs before or at startIndex.
   final reviews = await firestore.query(
-    from: 'reviews',
-    orderBy: orderBy('landed_index', false),
-    where: fieldLessThanOrEqual('landed_index', startIndex),
-    limit: 2,
+    StructuredQuery()
+      ..from = inCollection('reviews')
+      ..orderBy = [orderBy('landed_index', false)]
+      ..where = fieldLessThanOrEqual('landed_index', startIndex)
+      ..limit = 2,
   );
-  final firstReview = reviews.first;
-  index1 = firstReview.fields['landed_index']!.integerValue!;
-  review = firstReview.name.split('/').last;
-  final secondReview = reviews.last;
-  index2 = secondReview.fields['landed_index']!.integerValue!;
-  review2 = secondReview.name.split('/').last;
-  index3 = (int.parse(index2) - 1).toString();
-  index4 = (int.parse(index2) - 2).toString();
+  final firstReview = ReviewRecord(reviews.first);
+  index1 = firstReview.landedIndex!;
+  review = firstReview.review;
+  final secondReview = ReviewRecord(reviews.last);
+  index2 = secondReview.landedIndex!;
+  review2 = secondReview.review;
+  index3 = (index2 - 1).toString();
+  index4 = (index2 - 2).toString();
 
   final patchsets = await firestore.query(
-    from: 'patchsets',
+    StructuredQuery()
+      ..from = inCollection('patchsets')
+      ..orderBy = [orderBy('number', true)],
     parent: 'reviews/$review',
-    orderBy: orderBy('number', true),
   );
-  final patchsetFields = patchsets.last.fields;
-  lastPatchset = patchsetFields['number']!.integerValue!;
+  final patchsetRecord = PatchsetRecord(patchsets.last);
+  lastPatchset = patchsetRecord.number;
   lastPatchsetRef = 'refs/changes/$review/$lastPatchset';
-  patchsetGroup = patchsetFields['patchset_group']!.integerValue!;
+  patchsetGroup = patchsetRecord.patchsetGroup;
   patchsetGroupRef = 'refs/changes/$review/$patchsetGroup';
-  earlyPatchset = '1';
+  earlyPatchset = 1;
   earlyPatchsetRef = 'refs/changes/$review/$earlyPatchset';
   final patchsets2 = await firestore.query(
-    from: 'patchsets',
+    StructuredQuery()
+      ..from = inCollection('patchsets')
+      ..orderBy = [orderBy('number', true)],
     parent: 'reviews/$review2',
-    orderBy: orderBy('number', true),
   );
-  patchset2 = patchsets2.last.fields['number']!.integerValue!;
+  patchset2 = PatchsetRecord(patchsets2.last).number;
   patchset2Ref = 'refs/changes/$review/$patchset2';
 
   // Get commit hashes for the landed reviews, and for a commit before them
   var commits = {
-    for (final index in [index1, index2, index3, index4])
-      index: (await firestore.query(
-        from: 'commits',
-        where: fieldEquals('index', int.parse(index)),
-        limit: 1,
-      )).first.name.split('/').last,
+    for (final index in [index1, index2, int.parse(index3), int.parse(index4)])
+      index.toString(): (await firestore.query(
+        StructuredQuery()
+          ..from = inCollection('commits')
+          ..where = fieldEquals('index', index)
+          ..limit = 1,
+      )).first.name!.split('/').last,
   };
-  commit1 = commits[index1]!;
-  commit2 = commits[index2]!;
+  commit1 = commits[index1.toString()]!;
+  commit2 = commits[index2.toString()]!;
   commit3 = commits[index3]!;
   commit4 = commits[index4]!;
 }
@@ -258,11 +270,12 @@
     // patchsets, find one that does and set testCommitsStart to that index.
     expect(lastPatchset, isNot(patchsetGroup));
     expect(patchsetGroup, isNot(earlyPatchset));
-    expect(int.parse(index2), lessThan(int.parse(index1)));
+    expect(index2, lessThan(index1));
     // reviewWithComments should have some comments, to test linking
     final comments = await firestore.query(
-      from: 'comments',
-      where: fieldEquals('review', int.parse(reviewWithComments)),
+      StructuredQuery()
+        ..from = inCollection('comments')
+        ..where = fieldEquals('review', int.parse(reviewWithComments)),
     );
     expect(comments, isNotEmpty);
   });
@@ -273,10 +286,11 @@
     final change1 = makeTryChange('approvals', newFailure, lastPatchsetRef);
     await makeTryjob('approvals', change1).process([change1]);
     var documents = await firestore.query(
-      from: 'try_results',
-      where: fieldEquals('name', 'approvals_test'),
+      StructuredQuery()
+        ..from = inCollection('try_results')
+        ..where = fieldEquals('name', 'approvals_test'),
     );
-    await firestore.approveResult(documents.single.toDocument());
+    await firestore.approveResult(documents.single);
     final change2 = makeTryChange(
       'approvals',
       newFailure,
@@ -285,10 +299,11 @@
     );
     await makeTryjob('approvals', change2).process([change2]);
     documents = await firestore.query(
-      from: 'try_results',
-      where: fieldEquals('name', 'approvals_2_test'),
+      StructuredQuery()
+        ..from = inCollection('try_results')
+        ..where = fieldEquals('name', 'approvals_2_test'),
     );
-    await firestore.approveResult(documents.single.toDocument());
+    await firestore.approveResult(documents.single);
 
     final change3 = makeChange('approvals', newFailure, commit1, commit4);
     final change3a = ChangeRecord.fromMap({
@@ -343,29 +358,30 @@
       isTrue,
     );
     var commentsQuery = await firestore.query(
-      from: 'comments',
-      where: fieldEquals('review', int.parse(reviewWithComments)),
+      StructuredQuery()
+        ..from = inCollection('comments')
+        ..where = fieldEquals('review', int.parse(reviewWithComments)),
     );
     final landedIndex =
-        commentsQuery.first.fields[fBlamelistStartIndex]!.integerValue!;
+        commentsQuery.first.fields![fBlamelistStartIndex]!.integerValue!;
     for (final item in commentsQuery) {
-      final fields = item.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.toDocument(), [
+      await firestore.updateFields(item, [
         fBlamelistStartIndex,
         fBlamelistEndIndex,
       ]);
     }
-    var reviewDocument = await firestore.getDocument(
+    var reviewDocument = await firestore.getDocument<ReviewRecord>(
       '${firestore.documents}/reviews/$reviewWithComments',
     );
-    expect(reviewDocument.fields!['landed_index']!.integerValue, landedIndex);
-    reviewDocument.fields!.remove('landed_index');
-    await firestore.updateFields(reviewDocument, ['landed_index']);
+    expect(reviewDocument.landedIndex.toString(), landedIndex);
+    reviewDocument.doc.fields!.remove('landed_index');
+    await firestore.updateFields(reviewDocument.doc, ['landed_index']);
 
     await firestore.linkReviewToCommit(
       int.parse(reviewWithComments),
@@ -376,19 +392,20 @@
       int.parse(landedIndex),
     );
     commentsQuery = await firestore.query(
-      from: 'comments',
-      where: fieldEquals('review', int.parse(reviewWithComments)),
+      StructuredQuery()
+        ..from = inCollection('comments')
+        ..where = fieldEquals('review', int.parse(reviewWithComments)),
     );
     for (final item in commentsQuery) {
-      final fields = item.fields;
+      final fields = item.fields!;
       expect(fields[fBlamelistStartIndex]!.integerValue, landedIndex);
       expect(fields[fBlamelistEndIndex]!.integerValue, landedIndex);
       expect(fields[fReview]!.integerValue, reviewWithComments);
     }
-    reviewDocument = await firestore.getDocument(
+    reviewDocument = await firestore.getDocument<ReviewRecord>(
       '${firestore.documents}/reviews/$reviewWithComments',
     );
-    expect(reviewDocument.fields!['landed_index']!.integerValue, landedIndex);
+    expect(reviewDocument.landedIndex.toString(), landedIndex);
   });
 }
 
@@ -399,47 +416,48 @@
 }) async {
   final buildbucketId = 'bbID_$name';
   final buildDocuments = await firestore.query(
-    from: 'try_builds',
-    where: fieldEquals('buildbucket_id', buildbucketId),
+    StructuredQuery()
+      ..from = inCollection('try_builds')
+      ..where = fieldEquals('buildbucket_id', buildbucketId),
   );
   expect(buildDocuments.length, 1);
-  final document = buildDocuments.single;
-  expect(document.getBool('success'), success);
+  final record = TryBuildRecord(buildDocuments.single);
+  expect(record.success, success);
   if (truncated != null) {
-    expect(document.getBool('truncated'), truncated);
+    expect(record.truncated, truncated);
   } else {
-    expect(document.fields.containsKey('truncated'), isFalse);
+    expect(record.doc.fields!.containsKey('truncated'), isFalse);
   }
 }
 
-Future<void> checkBuild(String? builder, String index, {bool? success}) async {
-  final document = await firestore.getDocument(
+Future<void> checkBuild(String? builder, Object index, {bool? success}) async {
+  final record = await firestore.getDocument<BuildRecord>(
     '${firestore.documents}/builds/$builder:$index',
   );
-  expect(document.fields!['success']!.booleanValue, success);
+  expect(record.success, success);
 }
 
 Future<void> checkResult(
   ChangeRecord change,
-  String startIndex,
-  String endIndex,
+  Object startIndex,
+  Object endIndex,
   Map<String, dynamic> expected,
 ) async {
   expect([fConfigurations, fApproved], containsAll(expected.keys));
-  final resultName = await firestore.findResult(
-    change,
-    int.parse(startIndex),
-    int.parse(endIndex),
-  );
+  final start = startIndex is int
+      ? startIndex
+      : int.parse(startIndex as String);
+  final end = endIndex is int ? endIndex : int.parse(endIndex as String);
+  final resultName = await firestore.findResult(change, start, end);
   expect(resultName, isNotNull);
-  final resultDocument = await firestore.getDocument(resultName!);
-  final data = untagMap(resultDocument.fields!);
-  expect(data[fName], change.name);
-  expect(data[fBlamelistStartIndex], int.parse(startIndex));
-  expect(data[fBlamelistEndIndex], int.parse(endIndex));
+  final resultDocument = await firestore.getDocument<Document>(resultName!);
+  final data = ResultRecord(resultDocument);
+  expect(data.name, change.name);
+  expect(data.blamelistStartIndex, start);
+  expect(data.blamelistEndIndex, end);
   expect(
-    data[fConfigurations],
-    unorderedEquals(expected[fConfigurations] ?? data[fConfigurations]),
+    data.configurations,
+    unorderedEquals(expected[fConfigurations] ?? data.configurations),
   );
-  expect(data[fApproved], expected[fApproved] ?? data[fApproved]);
+  expect(data.approved, expected[fApproved] ?? data.approved);
 }
diff --git a/builder/test/builder_test.dart b/builder/test/builder_test.dart
index f670b6b..6d11d54 100644
--- a/builder/test/builder_test.dart
+++ b/builder/test/builder_test.dart
@@ -57,7 +57,7 @@
 
   for (final test in testsToRemove) {
     await deleteDocuments(
-      await firestore.runQuery(
+      await firestore.query(
         StructuredQuery()
           ..from = inCollection('results')
           ..where = fieldEquals(fName, test),
@@ -66,7 +66,7 @@
   }
   for (final builder in buildersToRemove) {
     await deleteDocuments(
-      await firestore.runQuery(
+      await firestore.query(
         StructuredQuery()
           ..from = inCollection('builds')
           ..where = fieldEquals('builder', builder),
diff --git a/builder/test/fakes.dart b/builder/test/fakes.dart
index 47a3630..9e90aa9 100644
--- a/builder/test/fakes.dart
+++ b/builder/test/fakes.dart
@@ -158,7 +158,7 @@
 
     result[fBlamelistEndIndex] = min<int>(endIndex, result[fBlamelistEndIndex]);
     if (!result[fConfigurations].contains(configuration)) {
-      result[fConfigurations] = List<String>.from(result[fConfigurations])
+      result[fConfigurations] = List<String?>.from(result[fConfigurations])
         ..add(configuration)
         ..sort();
     }
@@ -166,7 +166,7 @@
       result[fActive] = true;
       if (!result[fActiveConfigurations].contains(configuration)) {
         result[fActiveConfigurations] =
-            List<String>.from(result[fActiveConfigurations])
+            List<String?>.from(result[fActiveConfigurations])
               ..add(configuration)
               ..sort();
       }
@@ -204,24 +204,25 @@
   }
 
   @override
-  Future<List<SafeDocument>> tryApprovals(int review) async {
+  Future<List<TryResultRecord>> tryApprovals(int review) async {
     return fakeTryResults
         .where(
           (result) => result[fReview] == review && result[fApproved] == true,
         )
         .map(taggedMap)
         .map(
-          (fields) => SafeDocument(
+          (fields) => TryResultRecord(
             Document()
               ..fields = fields
-              ..name = '',
+              ..name =
+                  'projects/dummy/databases/(default)/documents/try_results/dummy',
           ),
         )
         .toList();
   }
 
   @override
-  Future<List<SafeDocument>> tryResults(
+  Future<List<TryResultRecord>> tryResults(
     int review,
     String configuration,
   ) async {
@@ -233,10 +234,11 @@
         )
         .map(taggedMap)
         .map(
-          (fields) => SafeDocument(
+          (fields) => TryResultRecord(
             Document()
               ..fields = fields
-              ..name = '',
+              ..name =
+                  'projects/dummy/databases/(default)/documents/try_results/dummy',
           ),
         )
         .toList();
diff --git a/builder/test/firestore_test.dart b/builder/test/firestore_test.dart
index dd073b8..ce56c16 100644
--- a/builder/test/firestore_test.dart
+++ b/builder/test/firestore_test.dart
@@ -55,7 +55,7 @@
 
     tearDown(() async {
       // Delete database records created by the tests.
-      var snapshot = await firestore.runQuery(
+      var snapshot = await firestore.query(
         StructuredQuery()
           ..from = inCollection('try_builds')
           ..where = fieldEquals('review', testReview),
@@ -64,14 +64,14 @@
         await firestore.deleteDocument(doc.name!);
       }
 
-      snapshot = await firestore.runQuery(
+      snapshot = await firestore.query(
         StructuredQuery()..from = inCollection('patchsets'),
         parent: 'reviews/$testReview/',
       );
       for (final doc in snapshot) {
         await firestore.deleteDocument(doc.name!);
       }
-      snapshot = await firestore.runQuery(
+      snapshot = await firestore.query(
         StructuredQuery()
           ..from = inCollection('results')
           ..where = fieldEquals('name', removeActiveConfigurationTestName),
@@ -117,7 +117,9 @@
         activeResult,
         'configuration 2',
       );
-      final document = await firestore.getDocument(createdResultDocument.name!);
+      final document = await firestore.getDocument<Document>(
+        createdResultDocument.name!,
+      );
       expect(document.fields, isNot(contains('active')));
       expect(document.fields, isNot(contains('active_configurations')));
       await firestore.deleteDocument(createdResultDocument.name!);
@@ -176,7 +178,7 @@
       });
       await firestore.storeTryChange(tryResult3, testReview, 3);
       // Set the results on patchsets 1 and 2 to approved.
-      final snapshot = await firestore.runQuery(
+      final snapshot = await firestore.query(
         StructuredQuery()
           ..from = inCollection('try_results')
           ..where = compositeFilter([
@@ -199,7 +201,7 @@
         'approved': true,
       }..remove('configuration');
       expect(1, approvals.length);
-      final approval = untagMap(approvals.single.fields);
+      final approval = untagMap(approvals.single.doc.fields!);
       expect(approval, expectedApproval);
     });
   });
diff --git a/builder/test/tryjob_test.dart b/builder/test/tryjob_test.dart
index e6b5f50..fc8ef60 100644
--- a/builder/test/tryjob_test.dart
+++ b/builder/test/tryjob_test.dart
@@ -35,30 +35,32 @@
 final testsToRemove = <String?>{};
 
 void registerChangeForDeletion(ChangeRecord change) {
-  testsToRemove.add(change.name);
   buildersToRemove.add(change.builderName);
+  testsToRemove.add(change.name);
 }
 
 Future<void> removeTryBuildersAndResults() async {
-  Future<void> deleteDocuments(List<SafeDocument> documents) async {
+  Future<void> deleteDocuments(List<Document> documents) async {
     for (final document in documents) {
-      await firestore.deleteDocument(document.name);
+      await firestore.deleteDocument(document.name!);
     }
   }
 
   for (final test in testsToRemove) {
     await deleteDocuments(
       await firestore.query(
-        from: 'try_results',
-        where: fieldEquals('name', test),
+        StructuredQuery()
+          ..from = inCollection('try_results')
+          ..where = fieldEquals('name', test),
       ),
     );
   }
   for (final builder in buildersToRemove) {
     await deleteDocuments(
       await firestore.query(
-        from: 'try_builds',
-        where: fieldEquals('builder', builder),
+        StructuredQuery()
+          ..from = inCollection('try_builds')
+          ..where = fieldEquals('builder', builder),
       ),
     );
   }
@@ -67,56 +69,59 @@
 Future<Map<String, String?>> loadTestCommits(int startIndex) async {
   // Get review data for the last two landed CLs before or at startIndex.
   final reviews = await firestore.query(
-    from: 'reviews',
-    orderBy: orderBy('landed_index', false),
-    where: fieldLessThanOrEqual('landed_index', startIndex),
-    limit: 2,
+    StructuredQuery()
+      ..from = inCollection('reviews')
+      ..orderBy = [orderBy('landed_index', false)]
+      ..where = fieldLessThanOrEqual('landed_index', startIndex)
+      ..limit = 2,
   );
-  final firstReview = reviews.first;
-  final String? index = firstReview.fields['landed_index']!.integerValue;
-  final String review = firstReview.name.split('/').last;
-  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)));
-  final String baseIndex = (int.parse(landedIndex) - 1).toString();
+  final firstReview = ReviewRecord(reviews.first);
+  final int index = firstReview.landedIndex!;
+  final String review = firstReview.review;
+  final secondReview = ReviewRecord(reviews.last);
+  final int landedIndex = secondReview.landedIndex!;
+  final String landedReview = secondReview.review;
+  final int baseIndex = landedIndex - 1;
 
   final patchsets = await firestore.query(
-    from: 'patchsets',
+    StructuredQuery()
+      ..from = inCollection('patchsets')
+      ..orderBy = [orderBy('number', true)],
     parent: 'reviews/$review',
-    orderBy: orderBy('number', true),
   );
-  final patchset = patchsets.last.fields['number']!.integerValue;
+  final patchset = PatchsetRecord(patchsets.last).number.toString();
   final previousPatchset = '1';
   final landedPatchsets = await firestore.query(
-    from: 'patchsets',
+    StructuredQuery()
+      ..from = inCollection('patchsets')
+      ..orderBy = [orderBy('number', true)],
     parent: 'reviews/$landedReview',
-    orderBy: orderBy('number', true),
   );
-  final landedPatchset = landedPatchsets.last.fields['number']!.integerValue;
+  final landedPatchset = PatchsetRecord(landedPatchsets.last).number.toString();
 
   // Get commit hashes for the landed reviews, and for a commit before them
   var commits = {
-    for (final index in [index, landedIndex, baseIndex])
-      index: (await firestore.query(
-        from: 'commits',
-        where: fieldEquals('index', int.parse(index!)),
-        limit: 1,
-      )).first.name.split('/').last,
+    for (final idx in [index, landedIndex, baseIndex])
+      idx: (await firestore.query(
+        StructuredQuery()
+          ..from = inCollection('commits')
+          ..where = fieldEquals('index', idx)
+          ..limit = 1,
+      )).first.name!.split('/').last,
   };
   return {
-    'index': index,
+    'index': index.toString(),
     'commit': commits[index],
     'review': review,
     'patchset': patchset,
     'patchsetRef': 'refs/changes/$review/$patchset',
     'previousPatchset': previousPatchset,
-    'landedIndex': landedIndex,
+    'landedIndex': landedIndex.toString(),
     'landedCommit': commits[landedIndex],
     'landedReview': landedReview,
     'landedPatchset': landedPatchset,
     'landedPatchsetRef': 'refs/changes/$landedReview/$landedPatchset',
-    'baseIndex': baseIndex,
+    'baseIndex': baseIndex.toString(),
     'baseCommit': commits[baseIndex],
   };
 }
@@ -184,15 +189,17 @@
 }) async {
   final buildbucketId = 'bbID_$name';
   final buildDocuments = await firestore.query(
-    from: 'try_builds',
-    where: fieldEquals('buildbucket_id', buildbucketId),
+    StructuredQuery()
+      ..from = inCollection('try_builds')
+      ..where = fieldEquals('buildbucket_id', buildbucketId),
   );
   expect(buildDocuments.length, 1);
-  expect(buildDocuments.single.fields['success']!.booleanValue, success);
+  final record = TryBuildRecord(buildDocuments.single);
+  expect(record.success, success);
   if (truncated != null) {
-    expect(buildDocuments.single.fields['truncated']!.booleanValue, truncated);
+    expect(record.truncated, truncated);
   } else {
-    expect(buildDocuments.single.fields.containsKey('truncated'), isFalse);
+    expect(record.doc.fields!.containsKey('truncated'), isFalse);
   }
 }
 
@@ -240,11 +247,12 @@
     expect(otherFailedStatus.success, isFalse);
     expect(otherFailedStatus.truncatedResults, isFalse);
     final result = await firestore.query(
-      from: 'try_results',
-      where: fieldEquals('name', 'failure_test'),
+      StructuredQuery()
+        ..from = inCollection('try_results')
+        ..where = fieldEquals('name', 'failure_test'),
     );
     expect(result.length, 1);
-    expect(result.single.getList('configurations')!.length, 2);
+    expect(TryResultRecord(result.single).configurations.length, 2);
   });
 
   test('landedFailure', () async {
@@ -282,7 +290,7 @@
   test('empty', () async {
     final emptyChange = makeChange('empty', 'Pass/Pass/Pass');
     final tryjob = makeTryjob('empty', emptyChange);
-    final status = await tryjob.process([]);
+    final status = await tryjob.process(<ChangeRecord>[]);
     await checkTryBuild('empty', success: true);
     expect(status.success, isTrue);
     expect(tryjob.success, isTrue);
@@ -310,13 +318,15 @@
     expect(tryjob.counter.hasTooManyFailingChanges, isFalse);
     expect(tryjob.counter.hasTruncatedChanges, isTrue);
     final existingResult = await firestore.query(
-      from: 'try_results',
-      where: fieldEquals('name', 'truncated_pass_2_test'),
+      StructuredQuery()
+        ..from = inCollection('try_results')
+        ..where = fieldEquals('name', 'truncated_pass_2_test'),
     );
     expect(existingResult.length, 1);
     final truncatedResult = await firestore.query(
-      from: 'try_results',
-      where: fieldEquals('name', 'truncatedPass_test'),
+      StructuredQuery()
+        ..from = inCollection('try_results')
+        ..where = fieldEquals('name', 'truncatedPass_test'),
     );
     expect(truncatedResult, isEmpty);
   });
@@ -342,19 +352,21 @@
     expect(tryjob.counter.hasTooManyFailingChanges, isTrue);
     expect(tryjob.counter.hasTruncatedChanges, isTrue);
     final existingResult = await firestore.query(
-      from: 'try_results',
-      where: fieldEquals('name', 'truncated_test'),
+      StructuredQuery()
+        ..from = inCollection('try_results')
+        ..where = fieldEquals('name', 'truncated_test'),
     );
     expect(existingResult.length, 1);
     final truncatedResult = await firestore.query(
-      from: 'try_results',
-      where: fieldEquals('name', 'truncated_2_test'),
+      StructuredQuery()
+        ..from = inCollection('try_results')
+        ..where = fieldEquals('name', 'truncated_2_test'),
     );
     expect(truncatedResult, isEmpty);
   });
 
   test('patchsets', () async {
-    final document = await firestore.getDocument(
+    final document = await firestore.getDocument<Document>(
       '${firestore.documents}/reviews/${data['review']}/patchsets/${data['patchset']}',
     );
     final fields = untagMap(document.fields!);
@@ -367,7 +379,7 @@
       fields['patchset_group'],
       fields['number'],
     );
-    final document1 = await firestore.getDocument(document.name!);
+    final document1 = await firestore.getDocument<Document>(document.name!);
     expect(untagMap(document1.fields!), equals(fields));
     fields['number'] += 1;
     fields['description'] = 'test description';
@@ -381,7 +393,7 @@
     );
     final name =
         '${firestore.documents}/reviews/${data['review']}/patchsets/${fields['number']}';
-    final document2 = await firestore.getDocument(name);
+    final document2 = await firestore.getDocument<Document>(name);
     final fields2 = untagMap(document2.fields!);
     expect(fields2, equals(fields));
     await firestore.deleteDocument(name);