1 files changed, 30 insertions, 15 deletions

diff --git a/document/src/test/java/com/yahoo/document/DocumentTypeTestCase.java b/document/src/test/java/com/yahoo/document/DocumentTypeTestCase.java
index abbfd31b778..47c7fe71343 100644
--- a/document/src/test/java/com/yahoo/document/DocumentTypeTestCase.java
+++ b/document/src/test/java/com/yahoo/document/DocumentTypeTestCase.java
@@ -3,10 +3,13 @@ package com.yahoo.document;
 
 import org.junit.Test;
 
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.Iterator;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
@@ -29,35 +32,49 @@ public class DocumentTypeTestCase {
     }
 
     @Test
+    public void testFieldSets() {
+        DocumentType root = new DocumentType("root");
+        root.addField("rootfield1", DataType.STRING);
+        root.addField("rootfield2", DataType.STRING);
+        root.addField("rootfield3", DataType.STRING);
+        root.addFieldSets(Collections.singletonMap(DocumentType.DOCUMENT, Arrays.asList("rootfield2")));
+        assertEquals(1, root.fieldSet().size());
+        assertEquals("rootfield2", root.fieldSet().iterator().next().getName());
+        assertEquals(3, root.fieldSetAll().size());
+    }
+
+    @Test
     public void testInheritance() {
         DocumentTypeManager typeManager = new DocumentTypeManager();
 
-        DocumentType child = new DocumentType("child");
-        Iterator inherited;
-
-        child.addField("childfield", DataType.INT);
-        child.addField("overridden", DataType.STRING);
+        DocumentType root = new DocumentType("root");
+        root.addField("rootfield", DataType.STRING);
+        root.addFieldSets(Collections.singletonMap(DocumentType.DOCUMENT, Arrays.asList("rootfield")));
 
         DocumentType parent1 = new DocumentType("parent1");
         parent1.addField("overridden", DataType.STRING);
         parent1.addField("parent1field", DataType.STRING);
-        child.inherit(parent1);
+        parent1.inherit(root);
+        parent1.addFieldSets(Collections.singletonMap(DocumentType.DOCUMENT, Arrays.asList("parent1field", "overridden")));
 
         DocumentType parent2 = new DocumentType("parent2");
         parent2.addField("parent2field", DataType.STRING);
-        child.inherit(parent2);
-
-        DocumentType root = new DocumentType("root");
-        root.addField("rootfield", DataType.STRING);
-        parent1.inherit(root);
         parent2.inherit(root);
+        parent2.addFieldSets(Collections.singletonMap(DocumentType.DOCUMENT, Arrays.asList("parent2field")));
+
+        DocumentType child = new DocumentType("child");
+        child.addField("childfield", DataType.INT);
+        child.addField("overridden", DataType.STRING);
+        child.inherit(parent1);
+        child.inherit(parent2);
+        child.addFieldSets(Collections.singletonMap(DocumentType.DOCUMENT, Arrays.asList("childfield", "overridden")));
 
         typeManager.register(root);
         typeManager.register(parent1);
         typeManager.register(parent2);
         typeManager.register(child);
 
-        inherited = child.getInheritedTypes().iterator();
+        Iterator inherited = child.getInheritedTypes().iterator();
         assertEquals(parent1, inherited.next());
         assertEquals(parent2, inherited.next());
         assertTrue(!inherited.hasNext());
@@ -95,9 +112,7 @@ public class DocumentTypeTestCase {
 
         assertFalse(fields.hasNext());
 
-        assert(child.getField("rootfield") != null);
-
-        // TODO: Test uninheriting
+        assertNotNull(child.getField("rootfield"));
     }
 
 }