2 files changed, 19 insertions, 13 deletions
diff --git a/Firestore/Example/Tests/Util/FSTHelpers.h b/Firestore/Example/Tests/Util/FSTHelpers.h
index d05c1f1..9b5f96a 100644
--- a/Firestore/Example/Tests/Util/FSTHelpers.h
+++ b/Firestore/Example/Tests/Util/FSTHelpers.h
@@ -245,8 +245,12 @@ FSTPatchMutation *FSTTestPatchMutation(
     NSDictionary<NSString *, id> *values,
     const std::vector<firebase::firestore::model::FieldPath> &updateMask);
 
-FSTTransformMutation *FSTTestTransformMutation(NSString *path,
-                                               NSArray<NSString *> *serverTimestampFields);
+/**
+ * Creates a FSTTransformMutation by parsing any FIRFieldValue sentinels in the provided data. The
+ * data is expected to use dotted-notation for nested fields (i.e.
+ * @{ @"foo.bar": [FIRFieldValue ...] } and must not contain any non-sentinel data.
+ */
+FSTTransformMutation *FSTTestTransformMutation(NSString *path, NSDictionary<NSString *, id> *data);
 
 /** Creates a delete mutation for the document key at the given path. */
 FSTDeleteMutation *FSTTestDeleteMutation(NSString *path);
diff --git a/Firestore/Example/Tests/Util/FSTHelpers.mm b/Firestore/Example/Tests/Util/FSTHelpers.mm
index 78b41a4..e1ed18c 100644
--- a/Firestore/Example/Tests/Util/FSTHelpers.mm
+++ b/Firestore/Example/Tests/Util/FSTHelpers.mm
@@ -122,7 +122,7 @@ NSDateComponents *FSTTestDateComponents(
   return comps;
 }
 
-FSTFieldValue *FSTTestFieldValue(id _Nullable value) {
+FSTUserDataConverter *FSTTestUserDataConverter() {
   // This owns the DatabaseIds since we do not have FirestoreClient instance to own them.
   static DatabaseId database_id{"project", DatabaseId::kDefault};
   FSTUserDataConverter *converter =
@@ -130,6 +130,11 @@ FSTFieldValue *FSTTestFieldValue(id _Nullable value) {
                                           preConverter:^id _Nullable(id _Nullable input) {
                                             return input;
                                           }];
+  return converter;
+}
+
+FSTFieldValue *FSTTestFieldValue(id _Nullable value) {
+  FSTUserDataConverter *converter = FSTTestUserDataConverter();
   // HACK: We use parsedQueryValue: since it accepts scalars as well as arrays / objects, and
   // our tests currently use FSTTestFieldValue() pretty generically so we don't know the intent.
   return [converter parsedQueryValue:value];
@@ -279,17 +284,14 @@ FSTPatchMutation *FSTTestPatchMutation(const absl::string_view path,
                                   precondition:Precondition::Exists(true)];
 }
 
-// For now this only creates TransformMutations with server timestamps.
-FSTTransformMutation *FSTTestTransformMutation(NSString *path,
-                                               NSArray<NSString *> *serverTimestampFields) {
+FSTTransformMutation *FSTTestTransformMutation(NSString *path, NSDictionary<NSString *, id> *data) {
   FSTDocumentKey *key = [FSTDocumentKey keyWithPath:testutil::Resource(util::MakeStringView(path))];
-  std::vector<FieldTransform> fieldTransforms;
-  for (NSString *field in serverTimestampFields) {
-    FieldPath fieldPath = testutil::Field(util::MakeStringView(field));
-    auto transformOp = absl::make_unique<ServerTimestampTransform>(ServerTimestampTransform::Get());
-    fieldTransforms.emplace_back(std::move(fieldPath), std::move(transformOp));
-  }
-  return [[FSTTransformMutation alloc] initWithKey:key fieldTransforms:std::move(fieldTransforms)];
+  FSTUserDataConverter *converter = FSTTestUserDataConverter();
+  FSTParsedUpdateData *result = [converter parsedUpdateData:data];
+  FSTCAssert(result.data.value.count == 0,
+             @"FSTTestTransformMutation() only expects transforms; no other data");
+  return [[FSTTransformMutation alloc] initWithKey:key
+                                   fieldTransforms:std::move(result.fieldTransforms)];
 }
 
 FSTDeleteMutation *FSTTestDeleteMutation(NSString *path) {