Added test cases for SortedMap

Fixed issue with custom comparators

hibernate-envers/src/main/java/org/hibernate/envers/configuration/metadata/CollectionMetadataGenerator.java

@@ -51,6 +51,8 @@ import org.hibernate.envers.entities.mapper.relation.BasicCollectionMapper;
 import org.hibernate.envers.entities.mapper.relation.CommonCollectionMapperData;
 import org.hibernate.envers.entities.mapper.relation.ListCollectionMapper;
 import org.hibernate.envers.entities.mapper.relation.MapCollectionMapper;
+import org.hibernate.envers.entities.mapper.relation.SortedMapCollectionMapper;
+import org.hibernate.envers.entities.mapper.relation.SortedSetCollectionMapper;
 import org.hibernate.envers.entities.mapper.relation.MiddleComponentData;
 import org.hibernate.envers.entities.mapper.relation.MiddleIdData;
 import org.hibernate.envers.entities.mapper.relation.ToOneIdMapper;
@@ -479,8 +481,8 @@ public final class CollectionMetadataGenerator {
         Type type = propertyValue.getType();
         if (type instanceof SortedSetType) {
 			currentMapper.addComposite(propertyAuditingData.getPropertyData(),
-                    new BasicCollectionMapper<Set>(commonCollectionMapperData,
+					new SortedSetCollectionMapper(commonCollectionMapperData,
-                    TreeSet.class, SortedSetProxy.class, elementComponentData));
+							TreeSet.class, SortedSetProxy.class, elementComponentData, propertyValue.getComparator()));
 		} else if (type instanceof SetType) {
 			currentMapper.addComposite(propertyAuditingData.getPropertyData(),
                     new BasicCollectionMapper<Set>(commonCollectionMapperData,
@@ -488,8 +490,8 @@ public final class CollectionMetadataGenerator {
         } else if (type instanceof SortedMapType) {
             // Indexed collection, so <code>indexComponentData</code> is not null.
 			currentMapper.addComposite(propertyAuditingData.getPropertyData(),
-                    new MapCollectionMapper<Map>(commonCollectionMapperData,
+					new SortedMapCollectionMapper(commonCollectionMapperData,
-                    TreeMap.class, SortedMapProxy.class, elementComponentData, indexComponentData));
+							TreeMap.class, SortedMapProxy.class, elementComponentData, indexComponentData, propertyValue.getComparator()));
         } else if (type instanceof MapType) {
             // Indexed collection, so <code>indexComponentData</code> is not null.
             currentMapper.addComposite(propertyAuditingData.getPropertyData(),

hibernate-envers/src/main/java/org/hibernate/envers/entities/mapper/relation/BasicCollectionMapper.java

@@ -38,8 +38,8 @@ import org.hibernate.collection.PersistentCollection;
 /**
  * @author Adam Warski (adam at warski dot org)
  */
-public final class BasicCollectionMapper<T extends Collection> extends AbstractCollectionMapper<T> implements PropertyMapper {
+public class BasicCollectionMapper<T extends Collection> extends AbstractCollectionMapper<T> implements PropertyMapper {
-    private final MiddleComponentData elementComponentData;
+    protected final MiddleComponentData elementComponentData;
 
     public BasicCollectionMapper(CommonCollectionMapperData commonCollectionMapperData,
                                  Class<? extends T> collectionClass, Class<? extends T> proxyClass,

hibernate-envers/src/main/java/org/hibernate/envers/entities/mapper/relation/MapCollectionMapper.java

@@ -38,9 +38,9 @@ import org.hibernate.collection.PersistentCollection;
 /**
  * @author Adam Warski (adam at warski dot org)
  */
-public final class MapCollectionMapper<T extends Map> extends AbstractCollectionMapper<T> implements PropertyMapper {
+public class MapCollectionMapper<T extends Map> extends AbstractCollectionMapper<T> implements PropertyMapper {
-    private final MiddleComponentData elementComponentData;
+    protected final MiddleComponentData elementComponentData;
-    private final MiddleComponentData indexComponentData;
+    protected final MiddleComponentData indexComponentData;
 
     public MapCollectionMapper(CommonCollectionMapperData commonCollectionMapperData,
                                Class<? extends T> collectionClass, Class<? extends T> proxyClass,

hibernate-envers/src/main/java/org/hibernate/envers/entities/mapper/relation/SortedMapCollectionMapper.java

@@ -0,0 +1,53 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2008, Red Hat Middleware LLC or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Middleware LLC.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.envers.entities.mapper.relation;
+
+import org.hibernate.envers.configuration.AuditConfiguration;
+import org.hibernate.envers.entities.mapper.relation.lazy.initializor.Initializor;
+import org.hibernate.envers.entities.mapper.relation.lazy.initializor.SortedMapCollectionInitializor;
+import org.hibernate.envers.reader.AuditReaderImplementor;
+
+import java.util.Comparator;
+import java.util.SortedMap;
+
+/**
+ * @author Michal Skowronek (mskowr at o2 dot pl)
+ */
+public final class SortedMapCollectionMapper extends MapCollectionMapper<SortedMap> {
+	private final Comparator comparator;
+
+	public SortedMapCollectionMapper(CommonCollectionMapperData commonCollectionMapperData,
+									 Class<? extends SortedMap> collectionClass, Class<? extends SortedMap> proxyClass,
+									 MiddleComponentData elementComponentData, MiddleComponentData indexComponentData, Comparator comparator) {
+		super(commonCollectionMapperData, collectionClass, proxyClass, elementComponentData, indexComponentData);
+		this.comparator = comparator;
+	}
+
+	protected Initializor<SortedMap> getInitializor(AuditConfiguration verCfg, AuditReaderImplementor versionsReader,
+													Object primaryKey, Number revision) {
+		return new SortedMapCollectionInitializor(verCfg, versionsReader, commonCollectionMapperData.getQueryGenerator(),
+				primaryKey, revision, collectionClass, elementComponentData, indexComponentData, comparator);
+	}
+
+}

hibernate-envers/src/main/java/org/hibernate/envers/entities/mapper/relation/SortedSetCollectionMapper.java

@@ -0,0 +1,53 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2008, Red Hat Middleware LLC or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Middleware LLC.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.envers.entities.mapper.relation;
+
+import org.hibernate.envers.configuration.AuditConfiguration;
+import org.hibernate.envers.entities.mapper.relation.lazy.initializor.Initializor;
+import org.hibernate.envers.entities.mapper.relation.lazy.initializor.SortedSetCollectionInitializor;
+import org.hibernate.envers.reader.AuditReaderImplementor;
+
+import java.util.Comparator;
+import java.util.SortedSet;
+
+/**
+ * @author Michal Skowronek (mskowr at o2 dot pl)
+ */
+public final class SortedSetCollectionMapper extends BasicCollectionMapper<SortedSet> {
+	private final Comparator comparator;
+
+	public SortedSetCollectionMapper(CommonCollectionMapperData commonCollectionMapperData,
+									 Class<? extends SortedSet> collectionClass, Class<? extends SortedSet> proxyClass,
+									 MiddleComponentData elementComponentData, Comparator comparator) {
+		super(commonCollectionMapperData, collectionClass, proxyClass, elementComponentData);
+		this.comparator = comparator;
+	}
+
+	protected Initializor<SortedSet> getInitializor(AuditConfiguration verCfg, AuditReaderImplementor versionsReader,
+													Object primaryKey, Number revision) {
+		return new SortedSetCollectionInitializor(verCfg, versionsReader, commonCollectionMapperData.getQueryGenerator(),
+				primaryKey, revision, collectionClass, elementComponentData, comparator);
+	}
+
+}

hibernate-envers/src/main/java/org/hibernate/envers/entities/mapper/relation/lazy/initializor/BasicCollectionInitializor.java

@@ -38,7 +38,7 @@ import org.hibernate.envers.reader.AuditReaderImplementor;
  * @author Adam Warski (adam at warski dot org)
  */
 public class BasicCollectionInitializor<T extends Collection> extends AbstractCollectionInitializor<T> {
-    private final Class<? extends T> collectionClass;
+    protected final Class<? extends T> collectionClass;
     private final MiddleComponentData elementComponentData;
 
     public BasicCollectionInitializor(AuditConfiguration verCfg,

hibernate-envers/src/main/java/org/hibernate/envers/entities/mapper/relation/lazy/initializor/MapCollectionInitializor.java

@@ -37,7 +37,7 @@ import org.hibernate.envers.reader.AuditReaderImplementor;
  * @author Adam Warski (adam at warski dot org)
  */
 public class MapCollectionInitializor<T extends Map> extends AbstractCollectionInitializor<T> {
-    private final Class<? extends T> collectionClass;
+    protected final Class<? extends T> collectionClass;
     private final MiddleComponentData elementComponentData;
     private final MiddleComponentData indexComponentData;
 

hibernate-envers/src/main/java/org/hibernate/envers/entities/mapper/relation/lazy/initializor/SortedMapCollectionInitializor.java

@@ -0,0 +1,72 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2008, Red Hat Middleware LLC or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Middleware LLC.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.envers.entities.mapper.relation.lazy.initializor;
+
+import org.hibernate.envers.configuration.AuditConfiguration;
+import org.hibernate.envers.entities.mapper.relation.MiddleComponentData;
+import org.hibernate.envers.entities.mapper.relation.query.RelationQueryGenerator;
+import org.hibernate.envers.exception.AuditException;
+import org.hibernate.envers.reader.AuditReaderImplementor;
+
+import java.lang.reflect.InvocationTargetException;
+import java.util.Comparator;
+import java.util.SortedMap;
+
+/**
+ * Initializes SortedMap collection with proper Comparator
+ *
+ * @author Michal Skowronek (mskowr at o2 dot pl)
+ */
+public class SortedMapCollectionInitializor extends MapCollectionInitializor<SortedMap> {
+	private final Comparator comparator;
+
+	public SortedMapCollectionInitializor(AuditConfiguration verCfg,
+										  AuditReaderImplementor versionsReader,
+										  RelationQueryGenerator queryGenerator,
+										  Object primaryKey, Number revision,
+										  Class<? extends SortedMap> collectionClass,
+										  MiddleComponentData elementComponentData,
+										  MiddleComponentData indexComponentData, Comparator comparator) {
+		super(verCfg, versionsReader, queryGenerator, primaryKey, revision, collectionClass, elementComponentData, indexComponentData);
+		this.comparator = comparator;
+	}
+
+	protected SortedMap initializeCollection(int size) {
+		if (comparator == null) {
+			return super.initializeCollection(size);
+		}
+		try {
+			return collectionClass.getConstructor(Comparator.class).newInstance(comparator);
+		} catch (InstantiationException e) {
+			throw new AuditException(e);
+		} catch (IllegalAccessException e) {
+			throw new AuditException(e);
+		} catch (NoSuchMethodException e) {
+			throw new AuditException(e);
+		} catch (InvocationTargetException e) {
+			throw new AuditException(e);
+		}
+	}
+
+}

hibernate-envers/src/main/java/org/hibernate/envers/entities/mapper/relation/lazy/initializor/SortedSetCollectionInitializor.java

@@ -0,0 +1,66 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2008, Red Hat Middleware LLC or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Middleware LLC.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.envers.entities.mapper.relation.lazy.initializor;
+
+import org.hibernate.envers.configuration.AuditConfiguration;
+import org.hibernate.envers.entities.mapper.relation.MiddleComponentData;
+import org.hibernate.envers.entities.mapper.relation.query.RelationQueryGenerator;
+import org.hibernate.envers.exception.AuditException;
+import org.hibernate.envers.reader.AuditReaderImplementor;
+
+import java.lang.reflect.InvocationTargetException;
+import java.util.Comparator;
+import java.util.SortedSet;
+
+/**
+ * Initializes SortedSet collection with proper Comparator
+ *
+ * @author Michal Skowronek (mskowr at o2 dot pl)
+ */
+public class SortedSetCollectionInitializor extends BasicCollectionInitializor<SortedSet> {
+	private final Comparator comparator;
+
+	public SortedSetCollectionInitializor(AuditConfiguration verCfg, AuditReaderImplementor versionsReader, RelationQueryGenerator queryGenerator, Object primaryKey, Number revision, Class<? extends SortedSet> collectionClass, MiddleComponentData elementComponentData, Comparator comparator) {
+		super(verCfg, versionsReader, queryGenerator, primaryKey, revision, collectionClass, elementComponentData);
+		this.comparator = comparator;
+	}
+
+	@Override
+	protected SortedSet initializeCollection(int size) {
+		if (comparator == null) {
+			return super.initializeCollection(size);
+		}
+		try {
+			return collectionClass.getConstructor(Comparator.class).newInstance(comparator);
+		} catch (InstantiationException e) {
+			throw new AuditException(e);
+		} catch (IllegalAccessException e) {
+			throw new AuditException(e);
+		} catch (NoSuchMethodException e) {
+			throw new AuditException(e);
+		} catch (InvocationTargetException e) {
+			throw new AuditException(e);
+		}
+	}
+}

hibernate-envers/src/test/java/org/hibernate/envers/test/entities/manytomany/SortedSetEntity.java

@@ -29,14 +29,14 @@ import org.hibernate.envers.Audited;
 import org.hibernate.envers.test.entities.StrTestEntity;
 import org.hibernate.envers.test.entities.StrTestEntityComparator;
 
-import javax.persistence.Entity;
+import javax.persistence.*;
-import javax.persistence.Id;
+import java.util.SortedMap;
-import javax.persistence.ManyToMany;
 import java.util.SortedSet;
+import java.util.TreeMap;
 import java.util.TreeSet;
 
 /**
- * Entity with custom-ordered SortedSet
+ * Entity with custom-ordered SortedSet and SortedMap
  *
  * @author Michal Skowronek (mskowr at o2 pl)
  */
@@ -52,6 +52,11 @@ public class SortedSetEntity {
     @ManyToMany
     @Sort(type = SortType.COMPARATOR, comparator = StrTestEntityComparator.class)
     private SortedSet<StrTestEntity> sortedSet = new TreeSet<StrTestEntity>(StrTestEntityComparator.INSTANCE);
+    @Audited
+	@ElementCollection
+	@MapKeyJoinColumn
+    @Sort(type = SortType.COMPARATOR, comparator = StrTestEntityComparator.class)
+    private SortedMap<StrTestEntity, String> sortedMap = new TreeMap<StrTestEntity, String>(StrTestEntityComparator.INSTANCE);
 
     public SortedSetEntity() {
     }
@@ -89,16 +94,21 @@ public class SortedSetEntity {
         this.sortedSet = sortedSet;
     }
 
+	public SortedMap<StrTestEntity, String> getSortedMap() {
+		return sortedMap;
+	}
+
+	public void setSortedMap(SortedMap<StrTestEntity, String> sortedMap) {
+		this.sortedMap = sortedMap;
+	}
+
 	public boolean equals(Object o) {
         if (this == o) return true;
         if (!(o instanceof SortedSetEntity)) return false;
 
         SortedSetEntity that = (SortedSetEntity) o;
 
-        if (data != null ? !data.equals(that.data) : that.data != null) return false;
+		return !(data != null ? !data.equals(that.data) : that.data != null) && !(id != null ? !id.equals(that.id) : that.id != null);
-        if (id != null ? !id.equals(that.id) : that.id != null) return false;
-
-        return true;
 	}
 
     public int hashCode() {

hibernate-envers/src/test/java/org/hibernate/envers/test/integration/manytomany/CustomComparatorEntityTest.java

@@ -29,20 +29,17 @@ import org.hibernate.envers.test.Priority;
 import org.hibernate.envers.test.entities.StrTestEntity;
 import org.hibernate.envers.test.entities.StrTestEntityComparator;
 import org.hibernate.envers.test.entities.manytomany.SortedSetEntity;
-import org.hibernate.testing.FailureExpected;
 import org.junit.Test;
 
 import javax.persistence.EntityManager;
-import java.util.Arrays;
+import java.util.*;
-import java.util.Iterator;
-import java.util.SortedSet;
 
 import static org.junit.Assert.assertEquals;
 
 /**
  * @author Michal Skowronek (mskowr at o2 pl)
  */
-public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest {
+public class CustomComparatorEntityTest extends AbstractEntityTest {
 
     private Integer id1;
     private Integer id2;
@@ -59,7 +56,7 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
     public void initData() {
         EntityManager em = getEntityManager();
 
-        SortedSetEntity entity1 = new SortedSetEntity(1, "sortedSet1");
+        SortedSetEntity entity1 = new SortedSetEntity(1, "sortedEntity1");
 
         // Revision 1
         em.getTransaction().begin();
@@ -77,6 +74,7 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
         em.persist(strTestEntity1);
         id1 = strTestEntity1.getId();
         entity1.getSortedSet().add(strTestEntity1);
+		entity1.getSortedMap().put(strTestEntity1, "abc");
 
         em.getTransaction().commit();
 
@@ -88,6 +86,7 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
         em.persist(strTestEntity2);
         id2 = strTestEntity2.getId();
         entity1.getSortedSet().add(strTestEntity2);
+		entity1.getSortedMap().put(strTestEntity2, "aaa");
 
         em.getTransaction().commit();
 
@@ -99,6 +98,7 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
         em.persist(strTestEntity3);
         id3 = strTestEntity3.getId();
         entity1.getSortedSet().add(strTestEntity3);
+		entity1.getSortedMap().put(strTestEntity3, "aba");
 
         em.getTransaction().commit();
 
@@ -110,6 +110,7 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
         em.persist(strTestEntity4);
         id4 = strTestEntity4.getId();
         entity1.getSortedSet().add(strTestEntity4);
+		entity1.getSortedMap().put(strTestEntity4, "aac");
 
         em.getTransaction().commit();
     }
@@ -127,7 +128,7 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
     public void testCurrentStateOfEntity1() {
         final SortedSetEntity entity1 = getEntityManager().find(SortedSetEntity.class, 1);
 
-        assertEquals("sortedSet1", entity1.getData());
+        assertEquals("sortedEntity1", entity1.getData());
         assertEquals(Integer.valueOf(1), entity1.getId());
 
         final SortedSet<StrTestEntity> sortedSet = entity1.getSortedSet();
@@ -138,6 +139,21 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
         checkStrTestEntity(iterator.next(), id4, "aac");
         checkStrTestEntity(iterator.next(), id3, "aba");
         checkStrTestEntity(iterator.next(), id1, "abc");
+
+		final SortedMap<StrTestEntity, String> sortedMap = entity1.getSortedMap();
+        assertEquals(StrTestEntityComparator.class, sortedMap.comparator().getClass());
+        assertEquals(4, sortedMap.size());
+		Iterator<Map.Entry<StrTestEntity, String>> mapIterator = sortedMap.entrySet().iterator();
+        checkStrTestEntity(mapIterator.next().getKey(), id2, "aaa");
+        checkStrTestEntity(mapIterator.next().getKey(), id4, "aac");
+        checkStrTestEntity(mapIterator.next().getKey(), id3, "aba");
+        checkStrTestEntity(mapIterator.next().getKey(), id1, "abc");
+
+	    mapIterator = sortedMap.entrySet().iterator();
+        assertEquals(mapIterator.next().getValue(), "aaa");
+        assertEquals(mapIterator.next().getValue(), "aac");
+        assertEquals(mapIterator.next().getValue(), "aba");
+        assertEquals(mapIterator.next().getValue(), "abc");
     }
 
     private void checkStrTestEntity(StrTestEntity entity, Integer id, String sortKey) {
@@ -146,20 +162,23 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
     }
 
     @Test
-    @FailureExpected(message = "Envers doesn't support custom comparators yet", jiraKey = "HHH-6176")
     public void testHistoryOfEntity1() throws Exception {
         SortedSetEntity entity1 = getAuditReader().find(SortedSetEntity.class, 1, 1);
 
-        assertEquals("sortedSet1", entity1.getData());
+        assertEquals("sortedEntity1", entity1.getData());
         assertEquals(Integer.valueOf(1), entity1.getId());
 
         SortedSet<StrTestEntity> sortedSet = entity1.getSortedSet();
         assertEquals(StrTestEntityComparator.class, sortedSet.comparator().getClass());
         assertEquals(0, sortedSet.size());
 
+		SortedMap<StrTestEntity, String> sortedMap = entity1.getSortedMap();
+        assertEquals(StrTestEntityComparator.class, sortedMap.comparator().getClass());
+        assertEquals(0, sortedMap.size());
+
         entity1 = getAuditReader().find(SortedSetEntity.class, 1, 2);
 
-        assertEquals("sortedSet1", entity1.getData());
+        assertEquals("sortedEntity1", entity1.getData());
         assertEquals(Integer.valueOf(1), entity1.getId());
 
         sortedSet = entity1.getSortedSet();
@@ -168,9 +187,18 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
         Iterator<StrTestEntity> iterator = sortedSet.iterator();
         checkStrTestEntity(iterator.next(), id1, "abc");
 
+		sortedMap = entity1.getSortedMap();
+        assertEquals(StrTestEntityComparator.class, sortedMap.comparator().getClass());
+        assertEquals(1, sortedMap.size());
+		Iterator<Map.Entry<StrTestEntity, String>> mapIterator = sortedMap.entrySet().iterator();
+        checkStrTestEntity(mapIterator.next().getKey(), id1, "abc");
+
+	    mapIterator = sortedMap.entrySet().iterator();
+        assertEquals(mapIterator.next().getValue(), "abc");
+
         entity1 = getAuditReader().find(SortedSetEntity.class, 1, 3);
 
-        assertEquals("sortedSet1", entity1.getData());
+        assertEquals("sortedEntity1", entity1.getData());
         assertEquals(Integer.valueOf(1), entity1.getId());
 
         sortedSet = entity1.getSortedSet();
@@ -180,9 +208,20 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
         checkStrTestEntity(iterator.next(), id2, "aaa");
         checkStrTestEntity(iterator.next(), id1, "abc");
 
+		sortedMap = entity1.getSortedMap();
+        assertEquals(StrTestEntityComparator.class, sortedMap.comparator().getClass());
+        assertEquals(2, sortedMap.size());
+		mapIterator = sortedMap.entrySet().iterator();
+        checkStrTestEntity(mapIterator.next().getKey(), id2, "aaa");
+        checkStrTestEntity(mapIterator.next().getKey(), id1, "abc");
+
+	    mapIterator = sortedMap.entrySet().iterator();
+        assertEquals(mapIterator.next().getValue(), "aaa");
+        assertEquals(mapIterator.next().getValue(), "abc");
+
         entity1 = getAuditReader().find(SortedSetEntity.class, 1, 4);
 
-        assertEquals("sortedSet1", entity1.getData());
+        assertEquals("sortedEntity1", entity1.getData());
         assertEquals(Integer.valueOf(1), entity1.getId());
 
         sortedSet = entity1.getSortedSet();
@@ -193,9 +232,22 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
         checkStrTestEntity(iterator.next(), id3, "aba");
         checkStrTestEntity(iterator.next(), id1, "abc");
 
+		sortedMap = entity1.getSortedMap();
+        assertEquals(StrTestEntityComparator.class, sortedMap.comparator().getClass());
+        assertEquals(3, sortedMap.size());
+		mapIterator = sortedMap.entrySet().iterator();
+        checkStrTestEntity(mapIterator.next().getKey(), id2, "aaa");
+        checkStrTestEntity(mapIterator.next().getKey(), id3, "aba");
+        checkStrTestEntity(mapIterator.next().getKey(), id1, "abc");
+
+	    mapIterator = sortedMap.entrySet().iterator();
+        assertEquals(mapIterator.next().getValue(), "aaa");
+        assertEquals(mapIterator.next().getValue(), "aba");
+        assertEquals(mapIterator.next().getValue(), "abc");
+
         entity1 = getAuditReader().find(SortedSetEntity.class, 1, 5);
 
-        assertEquals("sortedSet1", entity1.getData());
+        assertEquals("sortedEntity1", entity1.getData());
         assertEquals(Integer.valueOf(1), entity1.getId());
 
         sortedSet = entity1.getSortedSet();
@@ -206,6 +258,21 @@ public class SortedSetWithCustomComparatorEntityTest extends AbstractEntityTest
         checkStrTestEntity(iterator.next(), id4, "aac");
         checkStrTestEntity(iterator.next(), id3, "aba");
         checkStrTestEntity(iterator.next(), id1, "abc");
+
+		sortedMap = entity1.getSortedMap();
+        assertEquals(StrTestEntityComparator.class, sortedMap.comparator().getClass());
+        assertEquals(4, sortedMap.size());
+		mapIterator = sortedMap.entrySet().iterator();
+        checkStrTestEntity(mapIterator.next().getKey(), id2, "aaa");
+        checkStrTestEntity(mapIterator.next().getKey(), id4, "aac");
+        checkStrTestEntity(mapIterator.next().getKey(), id3, "aba");
+        checkStrTestEntity(mapIterator.next().getKey(), id1, "abc");
+
+	    mapIterator = sortedMap.entrySet().iterator();
+        assertEquals(mapIterator.next().getValue(), "aaa");
+        assertEquals(mapIterator.next().getValue(), "aac");
+        assertEquals(mapIterator.next().getValue(), "aba");
+        assertEquals(mapIterator.next().getValue(), "abc");
     }
 
 }