BAEL-1433: Article complete pushing final code.

core-java/src/test/java/com/baeldung/collection/WhenUsingTreeSet.java

@@ -1,8 +1,11 @@
 package com.baeldung.collection;
 
+import java.util.Comparator;
 import java.util.ConcurrentModificationException;
 import java.util.Iterator;
+import java.util.NavigableSet;
 import java.util.Set;
+import java.util.SortedSet;
 import java.util.TreeSet;
 
 import org.junit.Assert;
@@ -10,6 +13,29 @@ import org.junit.Test;
 
 public class WhenUsingTreeSet {
 
+    private static class Element {
+        private Integer id;
+
+        public Integer getId() {
+            return id;
+        }
+
+        public void setId(Integer id) {
+            this.id = id;
+        }
+
+        @Override
+        public String toString() {
+            // TODO Auto-generated method stub
+            return id.toString();
+        }
+    }
+
+    private Comparator<Element> comparator = (ele1, ele2) -> {
+        return ele1.getId()
+            .compareTo(ele2.getId());
+    };
+
     @Test
     public void whenAddingElement_shouldAddElement() {
         Set<String> treeSet = new TreeSet<>();
@@ -102,4 +128,96 @@ public class WhenUsingTreeSet {
         }
         Assert.assertEquals(2, treeSet.size());
     }
+
+    @Test(expected = NullPointerException.class)
+    public void whenAddingNullToNonEmptyTreeSet_shouldThrowException() {
+        Set<String> treeSet = new TreeSet<>();
+        treeSet.add("First");
+        treeSet.add(null);
+    }
+
+    @Test
+    public void whenCheckingFirstElement_shouldReturnFirstElement() {
+        TreeSet<String> treeSet = new TreeSet<>();
+        treeSet.add("First");
+        Assert.assertEquals("First", treeSet.first());
+    }
+
+    @Test
+    public void whenCheckingLastElement_shouldReturnLastElement() {
+        TreeSet<String> treeSet = new TreeSet<>();
+        treeSet.add("First");
+        treeSet.add("Last");
+        Assert.assertEquals("Last", treeSet.last());
+    }
+
+    @Test
+    public void whenUsingComparator_shouldSortAndInsertElements() {
+        Set<Element> treeSet = new TreeSet<>(comparator);
+        Element ele1 = new Element();
+        ele1.setId(100);
+        Element ele2 = new Element();
+        ele2.setId(200);
+
+        treeSet.add(ele1);
+        treeSet.add(ele2);
+
+        System.out.println(treeSet);
+    }
+
+    @Test
+    public void whenUsingHeadSet_shouldReturnElementsLessThanSpecifiedElement() {
+        Set<Element> treeSet = new TreeSet<>(comparator);
+        Element ele1 = new Element();
+        ele1.setId(100);
+        Element ele2 = new Element();
+        ele2.setId(200);
+
+        treeSet.add(ele1);
+        treeSet.add(ele2);
+
+        System.out.println(treeSet);
+    }
+
+    @Test
+    public void whenUsingSubSet_shouldReturnSubSetElements() {
+        SortedSet<Integer> treeSet = new TreeSet<>();
+        treeSet.add(1);
+        treeSet.add(2);
+        treeSet.add(3);
+        treeSet.add(4);
+        treeSet.add(5);
+        treeSet.add(6);
+
+        Set<Integer> subSet = treeSet.subSet(2, 6);
+        System.out.println(subSet);
+    }
+
+    @Test
+    public void whenUsingHeadSet_shouldReturnHeadSetElements() {
+        SortedSet<Integer> treeSet = new TreeSet<>();
+        treeSet.add(1);
+        treeSet.add(2);
+        treeSet.add(3);
+        treeSet.add(4);
+        treeSet.add(5);
+        treeSet.add(6);
+
+        Set<Integer> subSet = treeSet.headSet(6);
+        Assert.assertEquals(subSet, treeSet.subSet(1, 6));
+    }
+
+    @Test
+    public void whenUsingTailSet_shouldReturnTailSetElements() {
+        NavigableSet<Integer> treeSet = new TreeSet<>();
+        treeSet.add(1);
+        treeSet.add(2);
+        treeSet.add(3);
+        treeSet.add(4);
+        treeSet.add(5);
+        treeSet.add(6);
+
+        Set<Integer> subSet = treeSet.tailSet(3);
+        Assert.assertEquals(subSet, treeSet.subSet(3, true, 6, true));
+    }
 }