HHH-8855 HHH-8640 corrected EntityGraph loadplan strategy, support inheritance in SubGraphs

hibernate-core/src/main/java/org/hibernate/engine/query/spi/EntityGraphQueryHint.java

@@ -24,6 +24,7 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+
 import javax.persistence.AttributeNode;
 import javax.persistence.EntityGraph;
 import javax.persistence.Subgraph;
@@ -36,6 +37,7 @@ import org.hibernate.hql.internal.ast.tree.FromElement;
 import org.hibernate.hql.internal.ast.tree.FromElementFactory;
 import org.hibernate.persister.collection.QueryableCollection;
 import org.hibernate.sql.JoinType;
+import org.hibernate.type.CollectionType;
 import org.hibernate.type.EntityType;
 import org.hibernate.type.Type;
 
@@ -85,6 +87,8 @@ public class EntityGraphQueryHint {
 
 			final String attributeName = attributeNode.getAttributeName();
 			final String className = origin.getClassName();
+			// TODO: This is ignored by collection types and probably wrong for entity types.  Presumably it screws
+			// with inheritance.
 			final String role = className + "." + attributeName;
 			final String classAlias = origin.getClassAlias();
 			final String originTableAlias = origin.getTableAlias();
@@ -120,6 +124,7 @@ public class EntityGraphQueryHint {
 						);
 					}
 					else if ( propertyType.isCollectionType() ) {
+						CollectionType collectionType = (CollectionType) propertyType;
 						final String[] columns = origin.toColumns( originTableAlias, attributeName, false );
 
 						final FromElementFactory fromElementFactory = new FromElementFactory(
@@ -127,9 +132,9 @@ public class EntityGraphQueryHint {
 								attributeName, classAlias, columns, false
 						);
 						final QueryableCollection queryableCollection = walker.getSessionFactoryHelper()
-								.requireQueryableCollection( role );
+								.requireQueryableCollection( collectionType.getRole() );
 						fromElement = fromElementFactory.createCollection(
-								queryableCollection, role, JoinType.LEFT_OUTER_JOIN, true, false
+								queryableCollection, collectionType.getRole(), JoinType.LEFT_OUTER_JOIN, true, false
 						);
 					}
 				}

hibernate-core/src/main/java/org/hibernate/loader/plan/build/internal/AbstractEntityGraphVisitationStrategy.java

@@ -49,10 +49,8 @@ import org.hibernate.loader.plan.spi.LoadPlan;
 import org.hibernate.loader.plan.spi.Return;
 import org.hibernate.persister.walking.spi.AssociationAttributeDefinition;
 import org.hibernate.persister.walking.spi.AttributeDefinition;
-import org.hibernate.persister.walking.spi.CollectionDefinition;
 import org.hibernate.persister.walking.spi.CollectionElementDefinition;
 import org.hibernate.persister.walking.spi.CollectionIndexDefinition;
-import org.hibernate.persister.walking.spi.CompositionDefinition;
 import org.hibernate.persister.walking.spi.EntityDefinition;
 import org.hibernate.persister.walking.spi.WalkingException;
 import org.jboss.logging.Logger;
@@ -73,6 +71,7 @@ import org.jboss.logging.Logger;
  * it is not, then depends on which property is used to apply this entity graph.
  *
  * @author Strong Liu <stliu@hibernate.org>
+ * @author Brett Meyer
  */
 public abstract class AbstractEntityGraphVisitationStrategy
 		extends AbstractLoadPlanBuildingAssociationVisitationStrategy {
@@ -87,10 +86,14 @@ public abstract class AbstractEntityGraphVisitationStrategy
 	protected static final FetchStrategy DEFAULT_EAGER = new FetchStrategy( FetchTiming.IMMEDIATE, FetchStyle.JOIN );
 	protected static final FetchStrategy DEFAULT_LAZY = new FetchStrategy( FetchTiming.DELAYED, FetchStyle.SELECT );
 	protected final LoadQueryInfluencers loadQueryInfluencers;
-	protected final ArrayDeque<GraphNodeImplementor> graphStack = new ArrayDeque<GraphNodeImplementor>();
-	protected final ArrayDeque<AttributeNodeImplementor> attributeStack = new ArrayDeque<AttributeNodeImplementor>();
-	//the attribute nodes defined in the current graph node (entity graph or subgraph) we're working on
-	protected Map<String, AttributeNodeImplementor> attributeNodeImplementorMap = Collections.emptyMap();
+	// Queue containing entity/sub graphs to be visited.
+	private final ArrayDeque<GraphNodeImplementor> graphStack = new ArrayDeque<GraphNodeImplementor>();
+	// Queue containing attributes being visited, used eventually to determine the fetch strategy.
+	private final ArrayDeque<AttributeNodeImplementor> attributeStack = new ArrayDeque<AttributeNodeImplementor>();
+	// Queue of maps containing the current graph node's attributes.  Used for fast lookup, instead of iterating
+	// over graphStack.peekLast().attributeImplementorNodes().
+	private final ArrayDeque<Map<String, AttributeNodeImplementor>> attributeMapStack
+			= new ArrayDeque<Map<String, AttributeNodeImplementor>>();
 	private EntityReturn rootEntityReturn;
 	private final LockMode lockMode;
 
@@ -106,24 +109,21 @@ public abstract class AbstractEntityGraphVisitationStrategy
 	public void start() {
 		super.start();
 		graphStack.addLast( getRootEntityGraph() );
-		attributeNodeImplementorMap = buildAttributeNodeMap();
 	}
 
 	@Override
 	public void finish() {
 		super.finish();
 		graphStack.removeLast();
-		attributeNodeImplementorMap = Collections.emptyMap();
 		//applying a little internal stack checking
-		if ( !graphStack.isEmpty() || !attributeStack.isEmpty() || !attributeNodeImplementorMap.isEmpty() ) {
+		if ( !graphStack.isEmpty() || !attributeStack.isEmpty() || !attributeMapStack.isEmpty() ) {
 			throw new WalkingException( "Internal stack error" );
 		}
 	}
 
 	@Override
 	public void startingEntity(final EntityDefinition entityDefinition) {
-		//TODO check if the passed in entity definition is the same as the root entity graph (a.k.a they are came from same entity class)?
-		//this maybe the root entity graph or a sub graph.
+		attributeMapStack.addLast( buildAttributeNodeMap() );
 		super.startingEntity( entityDefinition );
 	}
 
@@ -143,6 +143,12 @@ public abstract class AbstractEntityGraphVisitationStrategy
 		return attributeNodeImplementorMap;
 	}
 
+	@Override
+	public void finishingEntity(final EntityDefinition entityDefinition) {
+		attributeMapStack.removeLast();
+		super.finishingEntity( entityDefinition );
+	}
+
 	/**
 	 * I'm using NULL-OBJECT pattern here, for attributes that not existing in the EntityGraph,
 	 * a predefined NULL-ATTRIBUTE-NODE is pushed to the stack.
@@ -156,12 +162,13 @@ public abstract class AbstractEntityGraphVisitationStrategy
 	 */
 	@Override
 	public boolean startingAttribute(AttributeDefinition attributeDefinition) {
+		Map<String, AttributeNodeImplementor> attributeMap = attributeMapStack.peekLast();
 		final String attrName = attributeDefinition.getName();
 		AttributeNodeImplementor attributeNode = NON_EXIST_ATTRIBUTE_NODE;
 		GraphNodeImplementor subGraphNode = NON_EXIST_SUBGRAPH_NODE;
 		//the attribute is in the EntityGraph, so, let's continue
-		if ( attributeNodeImplementorMap.containsKey( attrName ) ) {
-			attributeNode = attributeNodeImplementorMap.get( attrName );
+		if ( attributeMap.containsKey( attrName ) ) {
+			attributeNode = attributeMap.get( attrName );
 			//here we need to check if there is a subgraph (or sub key graph if it is an indexed attribute )
 			Map<Class, Subgraph> subGraphs = attributeNode.getSubgraphs();
 			Class javaType = attributeDefinition.getType().getReturnedClass();
@@ -183,45 +190,18 @@ public abstract class AbstractEntityGraphVisitationStrategy
 		super.finishingAttribute( attributeDefinition );
 	}
 
-	@Override
-	protected boolean handleAssociationAttribute(
-			final AssociationAttributeDefinition attributeDefinition) {
-		return super.handleAssociationAttribute( attributeDefinition );
-	}
-
-	@Override
-	protected boolean handleCompositeAttribute(
-			final AttributeDefinition attributeDefinition) {
-		return super.handleCompositeAttribute( attributeDefinition );
-	}
-
-
-	@Override
-	public void startingComposite(final CompositionDefinition compositionDefinition) {
-		super.startingComposite( compositionDefinition );
-	}
-
-
-	@Override
-	public void finishingComposite(final CompositionDefinition compositionDefinition) {
-		super.finishingComposite( compositionDefinition );
-	}
-
-
-	@Override
-	public void startingCollection(final CollectionDefinition collectionDefinition) {
-		super.startingCollection( collectionDefinition );
-	}
-
-	@Override
-	public void finishingCollection(final CollectionDefinition collectionDefinition) {
-		super.finishingCollection( collectionDefinition );
-	}
-
 
 	@Override
 	public void startingCollectionElements(
 			final CollectionElementDefinition elementDefinition) {
+		AttributeNodeImplementor attributeNode = attributeStack.peekLast();
+		GraphNodeImplementor subGraphNode = NON_EXIST_SUBGRAPH_NODE;
+		Map<Class, Subgraph> subGraphs = attributeNode.getSubgraphs();
+		Class javaType = elementDefinition.getType().getReturnedClass();
+		if ( !subGraphs.isEmpty() && subGraphs.containsKey( javaType ) ) {
+			subGraphNode = (GraphNodeImplementor) subGraphs.get( javaType );
+		}
+		graphStack.addLast( subGraphNode );
 		super.startingCollectionElements( elementDefinition );
 	}
 
@@ -229,6 +209,7 @@ public abstract class AbstractEntityGraphVisitationStrategy
 	public void finishingCollectionElements(
 			final CollectionElementDefinition elementDefinition) {
 		super.finishingCollectionElements( elementDefinition );
+		graphStack.removeLast();
 	}
 
 

hibernate-entitymanager/src/main/java/org/hibernate/jpa/graph/internal/SubgraphImpl.java

@@ -122,7 +122,7 @@ public class SubgraphImpl<T> extends AbstractGraphNode<T> implements Subgraph<T>
 	@Override
 	@SuppressWarnings("unchecked")
 	protected Attribute<T,?> resolveAttribute(String attributeName) {
-		final Attribute<T,?> attribute = managedType.getDeclaredAttribute( attributeName );
+		final Attribute<T,?> attribute = managedType.getAttribute( attributeName );
 		if ( attribute == null ) {
 			throw new IllegalArgumentException(
 					String.format(

hibernate-entitymanager/src/test/java/org/hibernate/jpa/test/graphs/Company.java

@@ -0,0 +1,54 @@
+/* 
+ * Hibernate, Relational Persistence for Idiomatic Java
+ * 
+ * JBoss, Home of Professional Open Source
+ * Copyright 2014 Red Hat Inc. and/or its affiliates and other contributors
+ * as indicated by the @authors tag. All rights reserved.
+ * See the copyright.txt in the distribution for a
+ * full listing of individual contributors.
+ *
+ * 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, v. 2.1.
+ * This program is distributed in the hope that it will be useful, but WITHOUT A
+ * 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,
+ * v.2.1 along with this distribution; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+ * MA  02110-1301, USA.
+ */
+package org.hibernate.jpa.test.graphs;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import javax.persistence.ElementCollection;
+import javax.persistence.Entity;
+import javax.persistence.FetchType;
+import javax.persistence.GeneratedValue;
+import javax.persistence.Id;
+import javax.persistence.OneToMany;
+import javax.persistence.OneToOne;
+
+
+/**
+ * @author Brett Meyer
+ */
+@Entity
+public class Company {
+	@Id @GeneratedValue
+	public long id;
+	
+	@OneToMany
+	public Set<Employee> employees = new HashSet<Employee>();
+	
+	@OneToOne(fetch = FetchType.LAZY)
+	public Location location;
+	
+	@ElementCollection
+	public Set<Market> markets = new HashSet<Market>();
+	
+	@ElementCollection(fetch = FetchType.EAGER)
+	public Set<String> phoneNumbers = new HashSet<String>();
+}

hibernate-entitymanager/src/test/java/org/hibernate/jpa/test/graphs/Employee.java

@@ -0,0 +1,48 @@
+/* 
+ * Hibernate, Relational Persistence for Idiomatic Java
+ * 
+ * JBoss, Home of Professional Open Source
+ * Copyright 2014 Red Hat Inc. and/or its affiliates and other contributors
+ * as indicated by the @authors tag. All rights reserved.
+ * See the copyright.txt in the distribution for a
+ * full listing of individual contributors.
+ *
+ * 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, v. 2.1.
+ * This program is distributed in the hope that it will be useful, but WITHOUT A
+ * 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,
+ * v.2.1 along with this distribution; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+ * MA  02110-1301, USA.
+ */
+package org.hibernate.jpa.test.graphs;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import javax.persistence.Entity;
+import javax.persistence.GeneratedValue;
+import javax.persistence.Id;
+import javax.persistence.Inheritance;
+import javax.persistence.InheritanceType;
+import javax.persistence.ManyToMany;
+
+
+/**
+ * @author Brett Meyer
+ */
+@Entity
+@Inheritance( strategy = InheritanceType.TABLE_PER_CLASS )
+public class Employee {
+	@Id @GeneratedValue
+	public long id;
+	
+	@ManyToMany
+	public Set<Manager> managers = new HashSet<Manager>();
+	
+	@ManyToMany
+	public Set<Employee> friends = new HashSet<Employee>();
+}

hibernate-entitymanager/src/test/java/org/hibernate/jpa/test/graphs/EntityGraphTest.java

@@ -21,8 +21,9 @@
  * 51 Franklin Street, Fifth Floor
  * Boston, MA  02110-1301  USA
  */
-package org.hibernate.jpa.test.graphs.find;
+package org.hibernate.jpa.test.graphs;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 import java.util.HashMap;
@@ -35,36 +36,38 @@ import javax.persistence.FetchType;
 import javax.persistence.GeneratedValue;
 import javax.persistence.Id;
 import javax.persistence.ManyToOne;
+import javax.persistence.Subgraph;
 
 import org.hibernate.Hibernate;
 import org.hibernate.jpa.test.BaseEntityManagerFunctionalTestCase;
+import org.hibernate.testing.TestForIssue;
 import org.junit.Test;
 
 /**
  * @author Christian Bauer
+ * @author Brett Meyer
  */
-public class FindEntityGraphTests extends BaseEntityManagerFunctionalTestCase {
+public class EntityGraphTest extends BaseEntityManagerFunctionalTestCase {
 
 	@Override
 	protected Class<?>[] getAnnotatedClasses() {
-		return new Class[] { Foo.class, Bar.class, Baz.class };
+		return new Class[] { Foo.class, Bar.class, Baz.class,
+				Company.class, Employee.class, Manager.class, Location.class };
 	}
 
 	@Test
+	@TestForIssue(jiraKey = "HHH-8857")
 	public void loadMultipleAssociations() {
 		EntityManager em = getOrCreateEntityManager();
 		em.getTransaction().begin();
 
 		Bar bar = new Bar();
-		bar.name = "bar";
 		em.persist( bar );
 
 		Baz baz = new Baz();
-		baz.name = "baz";
 		em.persist( baz );
 
 		Foo foo = new Foo();
-		foo.name = "foo";
 		foo.bar = bar;
 		foo.baz = baz;
 		em.persist( foo );
@@ -90,15 +93,65 @@ public class FindEntityGraphTests extends BaseEntityManagerFunctionalTestCase {
 		em.close();
 	}
 
+	/**
+	 * JPA 2.1 spec: "Add a node to the graph that corresponds to a managed type with inheritance. This allows for
+	 * multiple subclass subgraphs to be defined for this node of the entity graph. Subclass subgraphs will
+	 * automatically include the specified attributes of superclass subgraphs."
+	 */
+	@Test
+	@TestForIssue(jiraKey = "HHH-8640")
+	public void inheritanceTest() {
+		EntityManager em = getOrCreateEntityManager();
+		em.getTransaction().begin();
+
+		Manager manager = new Manager();
+		em.persist( manager );
+		Employee employee = new Employee();
+		employee.friends.add( manager );
+		employee.managers.add( manager );
+		em.persist( employee );
+		Company company = new Company();
+		company.employees.add( employee );
+		company.employees.add( manager );
+		em.persist( company );
+
+		em.getTransaction().commit();
+		em.clear();
+
+		em.getTransaction().begin();
+
+		EntityGraph<Company> entityGraph = em.createEntityGraph( Company.class );
+		Subgraph<Employee> subgraph = entityGraph.addSubgraph( "employees" );
+		subgraph.addAttributeNodes( "managers" );
+		subgraph.addAttributeNodes( "friends" );
+		Subgraph<Manager> subSubgraph = subgraph.addSubgraph( "managers", Manager.class );
+		subSubgraph.addAttributeNodes( "managers" );
+		subSubgraph.addAttributeNodes( "friends" );
+
+		Map<String, Object> properties = new HashMap<String, Object>();
+		properties.put( "javax.persistence.loadgraph", entityGraph );
+
+		Company result = em.find( Company.class, company.id, properties );
+
+		assertTrue( Hibernate.isInitialized( result ) );
+		assertTrue( Hibernate.isInitialized( result.employees ) );
+		assertEquals( result.employees.size(), 2 );
+		for (Employee resultEmployee : result.employees) {
+			assertTrue( Hibernate.isInitialized( resultEmployee.managers ) );
+			assertTrue( Hibernate.isInitialized( resultEmployee.friends ) );
+		}
+
+		em.getTransaction().commit();
+		em.close();
+	}
+
 	@Entity
-	public static class Foo {
+	private static class Foo {
 
 		@Id
 		@GeneratedValue
 		public Integer id;
 
-		public String name;
-
 		@ManyToOne(fetch = FetchType.LAZY)
 		public Bar bar;
 
@@ -107,23 +160,19 @@ public class FindEntityGraphTests extends BaseEntityManagerFunctionalTestCase {
 	}
 
 	@Entity
-	public static class Bar {
+	private static class Bar {
 
 		@Id
 		@GeneratedValue
-		public Integer id;
-
-		public String name;
+		private Integer id;
 	}
 
 	@Entity
-	public static class Baz {
+	private static class Baz {
 
 		@Id
 		@GeneratedValue
-		public Integer id;
-
-		public String name;
+		private Integer id;
 	}
 
 }

hibernate-entitymanager/src/test/java/org/hibernate/jpa/test/graphs/Location.java

@@ -0,0 +1,38 @@
+/* 
+ * Hibernate, Relational Persistence for Idiomatic Java
+ * 
+ * JBoss, Home of Professional Open Source
+ * Copyright 2014 Red Hat Inc. and/or its affiliates and other contributors
+ * as indicated by the @authors tag. All rights reserved.
+ * See the copyright.txt in the distribution for a
+ * full listing of individual contributors.
+ *
+ * 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, v. 2.1.
+ * This program is distributed in the hope that it will be useful, but WITHOUT A
+ * 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,
+ * v.2.1 along with this distribution; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+ * MA  02110-1301, USA.
+ */
+package org.hibernate.jpa.test.graphs;
+
+import javax.persistence.Entity;
+import javax.persistence.GeneratedValue;
+import javax.persistence.Id;
+
+/**
+ * @author Brett Meyer
+ */
+@Entity
+public class Location {
+	@Id @GeneratedValue
+	public long id;
+	
+	public String address;
+	
+	public int zip;
+}

hibernate-entitymanager/src/test/java/org/hibernate/jpa/test/graphs/Manager.java

@@ -0,0 +1,32 @@
+/* 
+ * Hibernate, Relational Persistence for Idiomatic Java
+ * 
+ * JBoss, Home of Professional Open Source
+ * Copyright 2014 Red Hat Inc. and/or its affiliates and other contributors
+ * as indicated by the @authors tag. All rights reserved.
+ * See the copyright.txt in the distribution for a
+ * full listing of individual contributors.
+ *
+ * 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, v. 2.1.
+ * This program is distributed in the hope that it will be useful, but WITHOUT A
+ * 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,
+ * v.2.1 along with this distribution; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+ * MA  02110-1301, USA.
+ */
+package org.hibernate.jpa.test.graphs;
+
+import javax.persistence.Entity;
+
+
+/**
+ * @author Brett Meyer
+ */
+@Entity
+public class Manager extends Employee {
+
+}

hibernate-entitymanager/src/test/java/org/hibernate/jpa/test/graphs/Market.java

@@ -0,0 +1,29 @@
+/* 
+ * Hibernate, Relational Persistence for Idiomatic Java
+ * 
+ * JBoss, Home of Professional Open Source
+ * Copyright 2014 Red Hat Inc. and/or its affiliates and other contributors
+ * as indicated by the @authors tag. All rights reserved.
+ * See the copyright.txt in the distribution for a
+ * full listing of individual contributors.
+ *
+ * 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, v. 2.1.
+ * This program is distributed in the hope that it will be useful, but WITHOUT A
+ * 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,
+ * v.2.1 along with this distribution; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+ * MA  02110-1301, USA.
+ */
+package org.hibernate.jpa.test.graphs;
+
+
+/**
+ * @author Brett Meyer
+ */
+public enum Market {
+	SERVICES, TECHNOLOGY, INDUSTRIAL;
+}

hibernate-entitymanager/src/test/java/org/hibernate/jpa/test/graphs/queryhint/QueryHintEntityGraphTest.java

@@ -23,27 +23,21 @@ package org.hibernate.jpa.test.graphs.queryhint;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
-import java.io.Serializable;
-import java.util.HashSet;
 import java.util.Iterator;
-import java.util.Set;
 
-import javax.persistence.ElementCollection;
-import javax.persistence.Entity;
 import javax.persistence.EntityGraph;
 import javax.persistence.EntityManager;
-import javax.persistence.FetchType;
-import javax.persistence.GeneratedValue;
-import javax.persistence.Id;
-import javax.persistence.ManyToMany;
-import javax.persistence.OneToMany;
-import javax.persistence.OneToOne;
 import javax.persistence.Query;
 import javax.persistence.Subgraph;
 
 import org.hibernate.Hibernate;
 import org.hibernate.jpa.QueryHints;
 import org.hibernate.jpa.test.BaseEntityManagerFunctionalTestCase;
+import org.hibernate.jpa.test.graphs.Company;
+import org.hibernate.jpa.test.graphs.Employee;
+import org.hibernate.jpa.test.graphs.Location;
+import org.hibernate.jpa.test.graphs.Manager;
+import org.hibernate.jpa.test.graphs.Market;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -81,10 +75,13 @@ public class QueryHintEntityGraphTest extends BaseEntityManagerFunctionalTestCas
 		entityManager = getOrCreateEntityManager();
 		entityManager.getTransaction().begin();
 		
-		Subgraph<Employee> subgraph1 = entityGraph.addSubgraph( "employees" );
-		subgraph1.addAttributeNodes( "managers" );
-		Subgraph<Employee> subgraph2 = subgraph1.addSubgraph( "managers" );
-		subgraph2.addAttributeNodes( "managers" );
+		Subgraph<Employee> subgraph = entityGraph.addSubgraph( "employees" );
+		subgraph.addAttributeNodes( "managers" );
+		subgraph.addAttributeNodes( "friends" );
+		Subgraph<Manager> subSubgraph = subgraph.addSubgraph( "managers", Manager.class );
+		subSubgraph.addAttributeNodes( "managers" );
+		subSubgraph.addAttributeNodes( "friends" );
+		
 		query = entityManager.createQuery( "from " + Company.class.getName() );
 		query.setHint( QueryHints.HINT_LOADGRAPH, entityGraph );
 		company = (Company) query.getSingleResult();
@@ -104,10 +101,14 @@ public class QueryHintEntityGraphTest extends BaseEntityManagerFunctionalTestCas
 		while (employeeItr.hasNext()) {
 			Employee employee = employeeItr.next();
 			assertTrue( Hibernate.isInitialized( employee.managers ) );
-			Iterator<Employee> managerItr =  employee.managers.iterator();
+			assertTrue( Hibernate.isInitialized( employee.friends ) );
+			// test 1 more level
+			Iterator<Manager> managerItr =  employee.managers.iterator();
 			while (managerItr.hasNext()) {
 				foundManager = true;
-				assertTrue( Hibernate.isInitialized( managerItr.next().managers ) );
+				Manager manager = managerItr.next();
+				assertTrue( Hibernate.isInitialized( manager.managers ) );
+				assertTrue( Hibernate.isInitialized( manager.friends ) );
 			}
 		}
 		assertTrue(foundManager);
@@ -144,16 +145,14 @@ public class QueryHintEntityGraphTest extends BaseEntityManagerFunctionalTestCas
 		EntityManager entityManager = getOrCreateEntityManager();
 		entityManager.getTransaction().begin();
 		
-		Employee manager1 = new Employee();
+		Manager manager1 = new Manager();
 		entityManager.persist( manager1 );
 		
-		Employee manager2 = new Employee();
-		manager2.managers = new HashSet<Employee>();
+		Manager manager2 = new Manager();
 		manager2.managers.add( manager1 );
 		entityManager.persist( manager2 );
 		
 		Employee employee = new Employee();
-		employee.managers = new HashSet<Employee>();
 		employee.managers.add( manager1 );
 		entityManager.persist( employee );
 		
@@ -163,15 +162,12 @@ public class QueryHintEntityGraphTest extends BaseEntityManagerFunctionalTestCas
 		entityManager.persist( location );
 		
 		Company company = new Company();
-		company.employees = new HashSet<Employee>();
 		company.employees.add( employee );
 		company.employees.add( manager1 );
 		company.employees.add( manager2 );
 		company.location = location;
-		company.markets = new HashSet<Market>();
 		company.markets.add( Market.SERVICES );
 		company.markets.add( Market.TECHNOLOGY );
-		company.phoneNumbers = new HashSet<String>();
 		company.phoneNumbers.add( "012-345-6789" );
 		company.phoneNumbers.add( "987-654-3210" );
 		entityManager.persist( company );
@@ -182,49 +178,6 @@ public class QueryHintEntityGraphTest extends BaseEntityManagerFunctionalTestCas
 
 	@Override
 	protected Class<?>[] getAnnotatedClasses() {
-		return new Class<?>[] { Company.class, Employee.class, Location.class };
-	}
-	
-	@Entity
-	private static class Company implements Serializable {
-		@Id @GeneratedValue
-		public long id;
-		
-		@OneToMany
-		public Set<Employee> employees;
-		
-		@OneToOne(fetch = FetchType.LAZY)
-		public Location location;
-		
-		@ElementCollection
-		public Set<Market> markets;
-		
-		@ElementCollection(fetch = FetchType.EAGER)
-		public Set<String> phoneNumbers;
-	}
-	
-	@Entity
-	private static class Employee {
-		@Id @GeneratedValue
-		public long id;
-		
-		@ManyToMany
-		public Set<Employee> managers;
-	}
-	
-	@Entity
-	private static class Location {
-		public Location() { }
-		
-		@Id @GeneratedValue
-		public long id;
-		
-		public String address;
-		
-		public int zip;
-	}
-	
-	private static enum Market {
-		SERVICES, TECHNOLOGY, INDUSTRIAL;
+		return new Class<?>[] { Company.class, Employee.class, Manager.class, Location.class };
 	}
 }