=== modified file 'dhis-2/dhis-api/src/main/java/org/hisp/dhis/preheat/Preheat.java'
--- dhis-2/dhis-api/src/main/java/org/hisp/dhis/preheat/Preheat.java	2016-02-01 07:12:51 +0000
+++ dhis-2/dhis-api/src/main/java/org/hisp/dhis/preheat/Preheat.java	2016-02-01 10:27:53 +0000
@@ -68,12 +68,12 @@
 
         if ( PreheatIdentifier.UID == identifier || PreheatIdentifier.AUTO == identifier )
         {
-            reference = get( identifier, object.getClass(), object.getUid() );
+            reference = get( PreheatIdentifier.UID, object.getClass(), object.getUid() );
         }
 
         if ( PreheatIdentifier.CODE == identifier || (reference == null && PreheatIdentifier.AUTO == identifier) )
         {
-            reference = get( identifier, object.getClass(), object.getCode() );
+            reference = get( PreheatIdentifier.CODE, object.getClass(), object.getCode() );
         }
 
         return reference;

=== modified file 'dhis-2/dhis-api/src/main/java/org/hisp/dhis/preheat/PreheatParams.java'
--- dhis-2/dhis-api/src/main/java/org/hisp/dhis/preheat/PreheatParams.java	2016-01-25 03:56:13 +0000
+++ dhis-2/dhis-api/src/main/java/org/hisp/dhis/preheat/PreheatParams.java	2016-02-01 10:27:53 +0000
@@ -46,7 +46,7 @@
 
     private Set<Class<? extends IdentifiableObject>> classes = new HashSet<>();
 
-    private Map<Class<? extends IdentifiableObject>, Set<String>> references = new HashMap<>();
+    private Map<PreheatIdentifier, Map<Class<? extends IdentifiableObject>, Set<String>>> references = new HashMap<>();
 
     public PreheatParams()
     {
@@ -85,12 +85,12 @@
         return this;
     }
 
-    public Map<Class<? extends IdentifiableObject>, Set<String>> getReferences()
+    public Map<PreheatIdentifier, Map<Class<? extends IdentifiableObject>, Set<String>>> getReferences()
     {
         return references;
     }
 
-    public PreheatParams setReferences( Map<Class<? extends IdentifiableObject>, Set<String>> references )
+    public PreheatParams setReferences( Map<PreheatIdentifier, Map<Class<? extends IdentifiableObject>, Set<String>>> references )
     {
         this.references = references;
         return this;

=== modified file 'dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/preheat/DefaultPreheatService.java'
--- dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/preheat/DefaultPreheatService.java	2016-02-01 09:41:43 +0000
+++ dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/preheat/DefaultPreheatService.java	2016-02-01 10:27:53 +0000
@@ -76,24 +76,40 @@
             for ( Class<? extends IdentifiableObject> klass : params.getClasses() )
             {
                 List<? extends IdentifiableObject> objects = manager.getAllNoAcl( klass ); // should we use getAll here? are we allowed to reference unshared objects?
-                preheat.put( params.getPreheatIdentifier(), objects );
+
+                if ( PreheatIdentifier.UID == params.getPreheatIdentifier() || PreheatIdentifier.AUTO == params.getPreheatIdentifier() )
+                {
+                    preheat.put( PreheatIdentifier.UID, objects );
+                }
+
+                if ( PreheatIdentifier.CODE == params.getPreheatIdentifier() || PreheatIdentifier.AUTO == params.getPreheatIdentifier() )
+                {
+                    preheat.put( PreheatIdentifier.CODE, objects );
+                }
             }
         }
         else if ( PreheatMode.REFERENCE == params.getPreheatMode() )
         {
-            for ( Class<? extends IdentifiableObject> klass : params.getReferences().keySet() )
+            Map<Class<? extends IdentifiableObject>, Set<String>> uidMap = params.getReferences().get( PreheatIdentifier.UID );
+            Map<Class<? extends IdentifiableObject>, Set<String>> codeMap = params.getReferences().get( PreheatIdentifier.CODE );
+
+            if ( uidMap != null && (PreheatIdentifier.UID == params.getPreheatIdentifier() || PreheatIdentifier.AUTO == params.getPreheatIdentifier()) )
             {
-                Collection<String> identifiers = params.getReferences().get( klass );
-
-                if ( PreheatIdentifier.UID == params.getPreheatIdentifier() )
+                for ( Class<? extends IdentifiableObject> klass : uidMap.keySet() )
                 {
+                    Collection<String> identifiers = uidMap.get( klass );
                     List<? extends IdentifiableObject> objects = manager.getByUid( klass, identifiers );
-                    preheat.put( params.getPreheatIdentifier(), objects );
+                    preheat.put( PreheatIdentifier.UID, objects );
                 }
-                else if ( PreheatIdentifier.CODE == params.getPreheatIdentifier() )
+            }
+
+            if ( codeMap != null && (PreheatIdentifier.CODE == params.getPreheatIdentifier() || PreheatIdentifier.AUTO == params.getPreheatIdentifier()) )
+            {
+                for ( Class<? extends IdentifiableObject> klass : codeMap.keySet() )
                 {
+                    Collection<String> identifiers = codeMap.get( klass );
                     List<? extends IdentifiableObject> objects = manager.getByCode( klass, identifiers );
-                    preheat.put( params.getPreheatIdentifier(), objects );
+                    preheat.put( PreheatIdentifier.CODE, objects );
                 }
             }
         }

=== modified file 'dhis-2/dhis-services/dhis-service-core/src/test/java/org/hisp/dhis/preheat/PreheatServiceTest.java'
--- dhis-2/dhis-services/dhis-service-core/src/test/java/org/hisp/dhis/preheat/PreheatServiceTest.java	2016-02-01 09:41:43 +0000
+++ dhis-2/dhis-services/dhis-service-core/src/test/java/org/hisp/dhis/preheat/PreheatServiceTest.java	2016-02-01 10:27:53 +0000
@@ -100,7 +100,8 @@
     public void testValidateRef()
     {
         PreheatParams params = new PreheatParams().setPreheatMode( PreheatMode.REFERENCE );
-        params.getReferences().put( DataElement.class, Sets.newHashSet( "ID1", "ID2" ) );
+        params.getReferences().put( PreheatIdentifier.UID, new HashMap<>() );
+        params.getReferences().get( PreheatIdentifier.UID ).put( DataElement.class, Sets.newHashSet( "ID1", "ID2" ) );
 
         preheatService.validate( params );
     }
@@ -362,10 +363,13 @@
         PreheatParams params = new PreheatParams();
         params.setPreheatMode( PreheatMode.REFERENCE );
 
-        Map<Class<? extends IdentifiableObject>, Set<String>> referenceMap = new HashMap<>();
-        referenceMap.put( DataElement.class, Sets.newHashSet( de1.getUid(), de2.getUid() ) );
-        referenceMap.put( User.class, Sets.newHashSet( user.getUid() ) );
-        params.setReferences( referenceMap );
+        Map<PreheatIdentifier, Map<Class<? extends IdentifiableObject>, Set<String>>> references = new HashMap<>();
+        references.put( PreheatIdentifier.UID, new HashMap<>() );
+
+        references.get( PreheatIdentifier.UID ).put( DataElement.class, Sets.newHashSet( de1.getUid(), de2.getUid() ) );
+        references.get( PreheatIdentifier.UID ).put( User.class, Sets.newHashSet( user.getUid() ) );
+
+        params.setReferences( references );
 
         preheatService.validate( params );
         Preheat preheat = preheatService.preheat( params );
@@ -412,9 +416,12 @@
         params.setPreheatIdentifier( PreheatIdentifier.CODE );
         params.setPreheatMode( PreheatMode.REFERENCE );
 
-        Map<Class<? extends IdentifiableObject>, Set<String>> references = new HashMap<>();
-        references.put( DataElement.class, Sets.newHashSet( de1.getCode(), de2.getCode() ) );
-        references.put( User.class, Sets.newHashSet( user.getCode() ) );
+        Map<PreheatIdentifier, Map<Class<? extends IdentifiableObject>, Set<String>>> references = new HashMap<>();
+        references.put( PreheatIdentifier.CODE, new HashMap<>() );
+
+        references.get( PreheatIdentifier.CODE ).put( DataElement.class, Sets.newHashSet( de1.getCode(), de2.getCode() ) );
+        references.get( PreheatIdentifier.CODE ).put( User.class, Sets.newHashSet( user.getCode() ) );
+
         params.setReferences( references );
 
         preheatService.validate( params );
@@ -440,12 +447,9 @@
         DataElementGroup dataElementGroup = fromJson( "preheat/degAUidRef.json", DataElementGroup.class );
         defaultSetup();
 
-        Map<Class<? extends IdentifiableObject>, Set<String>> references = preheatService.collectReferences( dataElementGroup )
-            .get( PreheatIdentifier.UID );
-
         PreheatParams params = new PreheatParams();
         params.setPreheatMode( PreheatMode.REFERENCE );
-        params.setReferences( references );
+        params.setReferences( preheatService.collectReferences( dataElementGroup ) );
 
         preheatService.validate( params );
         Preheat preheat = preheatService.preheat( params );
@@ -469,13 +473,10 @@
         DataElementGroup dataElementGroup = fromJson( "preheat/degACodeRef.json", DataElementGroup.class );
         defaultSetup();
 
-        Map<Class<? extends IdentifiableObject>, Set<String>> references = preheatService.collectReferences( dataElementGroup )
-            .get( PreheatIdentifier.CODE );
-
         PreheatParams params = new PreheatParams();
         params.setPreheatIdentifier( PreheatIdentifier.CODE );
         params.setPreheatMode( PreheatMode.REFERENCE );
-        params.setReferences( references );
+        params.setReferences( preheatService.collectReferences( dataElementGroup ) );
 
         preheatService.validate( params );
         Preheat preheat = preheatService.preheat( params );
@@ -499,12 +500,9 @@
         DataElementGroup dataElementGroup = fromJson( "preheat/degAUidRef.json", DataElementGroup.class );
         defaultSetup();
 
-        Map<Class<? extends IdentifiableObject>, Set<String>> references = preheatService.collectReferences( dataElementGroup )
-            .get( PreheatIdentifier.UID );
-
         PreheatParams params = new PreheatParams();
         params.setPreheatMode( PreheatMode.REFERENCE );
-        params.setReferences( references );
+        params.setReferences( preheatService.collectReferences( dataElementGroup ) );
 
         preheatService.validate( params );
         Preheat preheat = preheatService.preheat( params );
@@ -531,13 +529,10 @@
         DataElementGroup dataElementGroup = fromJson( "preheat/degACodeRef.json", DataElementGroup.class );
         defaultSetup();
 
-        Map<Class<? extends IdentifiableObject>, Set<String>> references = preheatService.collectReferences( dataElementGroup )
-            .get( PreheatIdentifier.CODE );
-
         PreheatParams params = new PreheatParams();
         params.setPreheatIdentifier( PreheatIdentifier.CODE );
         params.setPreheatMode( PreheatMode.REFERENCE );
-        params.setReferences( references );
+        params.setReferences( preheatService.collectReferences( dataElementGroup ) );
 
         preheatService.validate( params );
         Preheat preheat = preheatService.preheat( params );
@@ -559,23 +554,19 @@
 
     @Test
     @SuppressWarnings( "unchecked" )
-    @Ignore
     public void testPreheatReferenceConnectAUTO()
     {
         DataElementGroup dataElementGroup = fromJson( "preheat/degAAutoRef.json", DataElementGroup.class );
         defaultSetup();
 
-        Map<Class<? extends IdentifiableObject>, Set<String>> references = preheatService.collectReferences( dataElementGroup )
-            .get( PreheatIdentifier.UID );
-
         PreheatParams params = new PreheatParams();
-        params.setPreheatIdentifier( PreheatIdentifier.CODE );
+        params.setPreheatIdentifier( PreheatIdentifier.AUTO );
         params.setPreheatMode( PreheatMode.REFERENCE );
-        params.setReferences( references );
+        params.setReferences( preheatService.collectReferences( dataElementGroup ) );
 
         preheatService.validate( params );
         Preheat preheat = preheatService.preheat( params );
-        preheatService.connectReferences( dataElementGroup, preheat, PreheatIdentifier.CODE );
+        preheatService.connectReferences( dataElementGroup, preheat, PreheatIdentifier.AUTO );
 
         List<DataElement> members = new ArrayList<>( dataElementGroup.getMembers() );