1 files changed, 6 insertions, 7 deletions
diff --git a/src/test/scala/com/google/gimd/modification/ModificationTestCase.scala b/src/test/scala/com/google/gimd/modification/ModificationTestCase.scala
index d30f1cf..2f9e09c 100644
--- a/src/test/scala/com/google/gimd/modification/ModificationTestCase.scala
+++ b/src/test/scala/com/google/gimd/modification/ModificationTestCase.scala
@@ -19,6 +19,7 @@ import file.{FileType, File}
 import org.junit.Test
 import org.junit.Assert._
 import query._
+import UserType._
 
 final class ModificationTestCase {
 
@@ -26,19 +27,17 @@ final class ModificationTestCase {
 
   case class TreeNode(id: Int, name: String)
   object TreeNodeType extends UserType[TreeNode] {
-    def toUserObject(m: Message): TreeNode =
-      new TreeNode(
-        m.one("id").intField.value,
-        m.one("name").stringField.value
-      )
+    val id = FieldSpecOne("id", IntField, _.id)
+    val name = FieldSpecOne("name", StringField, _.name)
+    def fields = id :: name
     override def children = Seq(nestedMember)
-    def fields = List(FieldSpec("id", IntField, _.id), FieldSpec("name", StringField, _.name))
+    def toUserObject(m: Message) = new TreeNode(id(m), name(m))
   }
   object MockFileType extends FileType[TreeNode] {
     val pathPrefix = None
     val pathSuffix = None
     val userType = TreeNodeType
-    def name(m: Message) = m.one("id").stringField.value
+    def name(m: Message) = userType.id(m).toString
   }
   case class MockFile(message: Message) extends File[TreeNode] {
     val path = ""