=== modified file 'dhis-2/dhis-api/src/main/java/org/hisp/dhis/program/Program.java'
--- dhis-2/dhis-api/src/main/java/org/hisp/dhis/program/Program.java	2015-09-11 11:54:53 +0000
+++ dhis-2/dhis-api/src/main/java/org/hisp/dhis/program/Program.java	2015-09-13 14:17:27 +0000
@@ -153,6 +153,34 @@
     // Logic methods
     // -------------------------------------------------------------------------
 
+    public void addOrganisationUnit( OrganisationUnit organisationUnit )
+    {
+        organisationUnits.add( organisationUnit );
+        organisationUnit.getPrograms().add( this );
+    }
+    
+    public void removeOrganisationUnit( OrganisationUnit organisationUnit )
+    {
+        organisationUnits.remove( organisationUnit );
+        organisationUnit.getPrograms().remove( this );
+    }
+    
+    public void updateOrganisationUnits( Set<OrganisationUnit> updates )
+    {
+        for ( OrganisationUnit unit : new HashSet<>( organisationUnits ) )
+        {
+            if ( !updates.contains( unit ) )
+            {
+                removeOrganisationUnit( unit );
+            }
+        }
+        
+        for ( OrganisationUnit unit : updates )
+        {
+            addOrganisationUnit( unit );
+        }
+    }
+    
     /**
      * Returns the ProgramTrackedEntityAttribute of this Program which contains
      * the given TrackedEntityAttribute.

=== modified file 'dhis-2/dhis-api/src/main/java/org/hisp/dhis/program/ProgramService.java'
--- dhis-2/dhis-api/src/main/java/org/hisp/dhis/program/ProgramService.java	2015-07-02 07:05:55 +0000
+++ dhis-2/dhis-api/src/main/java/org/hisp/dhis/program/ProgramService.java	2015-09-13 14:17:27 +0000
@@ -247,6 +247,8 @@
      */
     List<Program> getByCurrentUser( ProgramType type );
 
+    void mergeWithCurrentUserOrganisationUnits( Program program, Collection<OrganisationUnit> mergeOrganisationUnits );
+    
     /**
      * @param htmlCode
      * @param program

=== modified file 'dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataset/DefaultDataSetService.java'
--- dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataset/DefaultDataSetService.java	2015-09-13 13:48:45 +0000
+++ dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataset/DefaultDataSetService.java	2015-09-13 14:17:27 +0000
@@ -422,17 +422,17 @@
     @Override
     public void mergeWithCurrentUserOrganisationUnits( DataSet dataSet, Collection<OrganisationUnit> mergeOrganisationUnits )
     {
-        Set<OrganisationUnit> dataSetOrgUnits = new HashSet<>( dataSet.getSources() );
+        Set<OrganisationUnit> selectedOrgUnits = new HashSet<>( dataSet.getSources() );
         
         OrganisationUnitQueryParams params = new OrganisationUnitQueryParams();
         params.setParents( currentUserService.getCurrentUser().getOrganisationUnits() );
 
         List<OrganisationUnit> userOrganisationUnits = organisationUnitService.getOrganisationUnitsByQuery( params );
 
-        dataSetOrgUnits.removeAll( userOrganisationUnits );
-        dataSetOrgUnits.addAll( mergeOrganisationUnits );
+        selectedOrgUnits.removeAll( userOrganisationUnits );
+        selectedOrgUnits.addAll( mergeOrganisationUnits );
 
-        dataSet.updateOrganisationUnits( dataSetOrgUnits );
+        dataSet.updateOrganisationUnits( selectedOrgUnits );
 
         updateDataSet( dataSet );
     }

=== modified file 'dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/program/DefaultProgramService.java'
--- dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/program/DefaultProgramService.java	2015-09-03 08:15:09 +0000
+++ dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/program/DefaultProgramService.java	2015-09-13 14:17:27 +0000
@@ -29,6 +29,7 @@
  */
 
 import com.google.common.collect.Sets;
+
 import org.apache.commons.lang3.StringUtils;
 import org.hisp.dhis.common.ValueType;
 import org.hisp.dhis.i18n.I18n;
@@ -36,6 +37,8 @@
 import org.hisp.dhis.i18n.I18nService;
 import org.hisp.dhis.option.Option;
 import org.hisp.dhis.organisationunit.OrganisationUnit;
+import org.hisp.dhis.organisationunit.OrganisationUnitQueryParams;
+import org.hisp.dhis.organisationunit.OrganisationUnitService;
 import org.hisp.dhis.trackedentity.TrackedEntity;
 import org.hisp.dhis.trackedentity.TrackedEntityAttribute;
 import org.hisp.dhis.trackedentity.TrackedEntityAttributeService;
@@ -53,6 +56,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Date;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.regex.Matcher;
@@ -107,6 +111,13 @@
     {
         this.userService = userService;
     }
+    
+    private OrganisationUnitService organisationUnitService;
+
+    public void setOrganisationUnitService( OrganisationUnitService organisationUnitService )
+    {
+        this.organisationUnitService = organisationUnitService;
+    }
 
     @Autowired
     private TrackedEntityAttributeService attributeService;
@@ -308,6 +319,23 @@
         return programs;
     }
 
+    @Override
+    public void mergeWithCurrentUserOrganisationUnits( Program program, Collection<OrganisationUnit> mergeOrganisationUnits )
+    {
+        Set<OrganisationUnit> selectedOrgUnits = new HashSet<>( program.getOrganisationUnits() );
+
+        OrganisationUnitQueryParams params = new OrganisationUnitQueryParams();
+        params.setParents( currentUserService.getCurrentUser().getOrganisationUnits() );
+
+        List<OrganisationUnit> userOrganisationUnits = organisationUnitService.getOrganisationUnitsByQuery( params );
+
+        selectedOrgUnits.removeAll( userOrganisationUnits );
+        selectedOrgUnits.addAll( mergeOrganisationUnits );
+        
+        program.updateOrganisationUnits( selectedOrgUnits );
+        
+        updateProgram( program );
+    }
 
     @Override
     public String prepareDataEntryFormForAdd( String htmlCode, Program program, Collection<User> healthWorkers,

=== modified file 'dhis-2/dhis-services/dhis-service-core/src/main/resources/META-INF/dhis/beans.xml'
--- dhis-2/dhis-services/dhis-service-core/src/main/resources/META-INF/dhis/beans.xml	2015-08-31 14:41:00 +0000
+++ dhis-2/dhis-services/dhis-service-core/src/main/resources/META-INF/dhis/beans.xml	2015-09-13 14:17:27 +0000
@@ -1134,6 +1134,7 @@
     <property name="i18nService" ref="org.hisp.dhis.i18n.I18nService" />
     <property name="userService" ref="org.hisp.dhis.user.UserService" />
     <property name="currentUserService" ref="org.hisp.dhis.user.CurrentUserService" />
+    <property name="organisationUnitService" ref="org.hisp.dhis.organisationunit.OrganisationUnitService" />
   </bean>
 
   <bean id="org.hisp.dhis.program.ProgramValidationService" class="org.hisp.dhis.program.DefaultProgramValidationService">

=== modified file 'dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-program/src/main/java/org/hisp/dhis/trackedentity/action/program/DefineProgramAssociationsAction.java'
--- dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-program/src/main/java/org/hisp/dhis/trackedentity/action/program/DefineProgramAssociationsAction.java	2015-01-17 07:41:26 +0000
+++ dhis-2/dhis-web/dhis-web-maintenance/dhis-web-maintenance-program/src/main/java/org/hisp/dhis/trackedentity/action/program/DefineProgramAssociationsAction.java	2015-09-13 14:17:27 +0000
@@ -29,8 +29,6 @@
  */
 
 import java.util.Collection;
-import java.util.HashSet;
-import java.util.Set;
 
 import org.hisp.dhis.organisationunit.OrganisationUnit;
 import org.hisp.dhis.oust.manager.SelectionTreeManager;
@@ -40,9 +38,7 @@
 import com.opensymphony.xwork2.Action;
 
 /**
- * @author Kristian
- * @version $Id: DefineDataSetAssociationsAction.java 3648 2007-10-15 22:47:45Z
- *          larshelg $
+ * @author Lars Helge Overland
  */
 public class DefineProgramAssociationsAction
     implements Action
@@ -84,49 +80,12 @@
     public String execute()
         throws Exception
     {
-        Collection<OrganisationUnit> rootUnits = selectionTreeManager.getRootOrganisationUnits();
-
-        Set<OrganisationUnit> unitsInTheTree = new HashSet<>();
-
-        getUnitsInTheTree( rootUnits, unitsInTheTree );
-
         Program program = programService.getProgram( id );
 
-        Set<OrganisationUnit> assignedUnits = program.getOrganisationUnits();
-
-        assignedUnits.removeAll( convert( unitsInTheTree ) );
-
-        Collection<OrganisationUnit> selectedOrganisationUnits = selectionTreeManager
-            .getReloadedSelectedOrganisationUnits();
-
-        assignedUnits.addAll( convert( selectedOrganisationUnits ) );
-
-        program.setOrganisationUnits( assignedUnits );
-
-        programService.updateProgram( program );
+        Collection<OrganisationUnit> units = selectionTreeManager.getReloadedSelectedOrganisationUnits();
+        
+        programService.mergeWithCurrentUserOrganisationUnits( program, units );
 
         return SUCCESS;
     }
-
-    // -------------------------------------------------------------------------
-    // Supporting methods
-    // -------------------------------------------------------------------------
-
-    private Set<OrganisationUnit> convert( Collection<OrganisationUnit> organisationUnits )
-    {
-        Set<OrganisationUnit> units = new HashSet<>();
-
-        units.addAll( organisationUnits );
-
-        return units;
-    }
-
-    private void getUnitsInTheTree( Collection<OrganisationUnit> rootUnits, Set<OrganisationUnit> unitsInTheTree )
-    {
-        for ( OrganisationUnit root : rootUnits )
-        {
-            unitsInTheTree.add( root );
-            getUnitsInTheTree( root.getChildren(), unitsInTheTree );
-        }
-    }
 }