Modifications to registerDirty( SObject, List<SObjectField> ) to only register dirtyFields

fflib/src/classes/fflib_ISObjectUnitOfWork.cls

@@ -60,13 +60,16 @@ public interface fflib_ISObjectUnitOfWork
      **/
     void registerDirty(SObject record);
     /**
-     * Register specific fields on record to be updated when work is commited
+     * Register specific fields on an existing record to be updated when work is committed.
      *
-     * If the record has previously been registered as dirty, the dirty fields on the record in this call will overwrite
-     * the values of the previously registered dirty record
+     * When called after registerDirty(record), only the specified fields on existing record get updated.
+     * When called first time, a new clean instance is made and only the specified fields are populated.
+     * When called followed by registerDirty(record), an exception is thrown, due to mismatch in memory reference.
+     *  NOTE: no reference is kept in this function. When applying registerDirty( record, fields ), then altering the input record,
+     *      one will have to call registerDirty( fields ) again to make sure those changes are included in the commitWork
      *
      * @param record An existing record
-     * @param dirtyFields The fields to update if record is already registered
+     * @param dirtyFields The list of fields which should only be registered
      **/
     void registerDirty(SObject record, List<SObjectField> dirtyFields);
     /**

fflib/src/classes/fflib_SObjectUnitOfWork.cls

@@ -259,22 +259,38 @@ public virtual class fflib_SObjectUnitOfWork
         if(!m_dirtyMapByType.containsKey(sObjectType))
             throw new UnitOfWorkException(String.format('SObject type {0} is not supported by this unit of work', new String[] { sObjectType }));
 
-        // If record isn't registered as dirty, or no dirty fields to drive a merge
-        if (!m_dirtyMapByType.get(sObjectType).containsKey(record.Id) || dirtyFields.isEmpty())
+        // If record isn't registered as dirty, register it
+        if (!m_dirtyMapByType.get(sObjectType).containsKey(record.Id))
         {
-            // Register the record as dirty
-            m_dirtyMapByType.get(sObjectType).put(record.Id, record);
+            if (dirtyFields.isEmpty() )
+            {
+                // When no dirty fields are specified, register full record dirty and end method
+                m_dirtyMapByType.get(sObjectType).put(record.Id, record);
+                return;
+            }
+            else
+            {
+                // When dirty fields are specified, register clean SObject to make sure only dirty fields get populated
+                m_dirtyMapByType.get(sObjectType).put(record.Id, record.getSObjectType().newSObject( record.Id ) );
+            }
+
         }
-        else
-        {
-            // Update the registered record's fields
-            SObject registeredRecord = m_dirtyMapByType.get(sObjectType).get(record.Id);
 
-            for (SObjectField dirtyField : dirtyFields) {
-                registeredRecord.put(dirtyField, record.get(dirtyField));
-            }
+        // Fetch existing or newly instantiated record to update the specified fields
+        SObject registeredRecord = m_dirtyMapByType.get(sObjectType).get(record.Id);
 
-            m_dirtyMapByType.get(sObjectType).put(record.Id, registeredRecord);
+        // When provided instance is different than registered one and NO fields are provided, system cannot determine
+        // what changes to apply (assuming updating nothing is incorrect)
+        // Scenario: first registerDirty(record, fields), followed by registerDirty( record ) 
+        if (registeredRecord !== record && dirtyFields.isEmpty())
+        {
+            throw new UnitOfWorkException('Cannot determine what fields to update on record ' + record);
+        }
+
+		// When records have same memory reference, or only a subset of fields are registered, update the registered record
+        for (SObjectField dirtyField : dirtyFields)
+        {
+            registeredRecord.put(dirtyField, record.get(dirtyField));
         }
     }
 

fflib/src/classes/fflib_SObjectUnitOfWorkTest.cls

@@ -386,6 +386,31 @@ private with sharing class fflib_SObjectUnitOfWorkTest
         System.assertEquals(opp.Amount, amountUpdate.Amount);
     }
 
+    /**
+     * Try registering a single field as dirty, while entry record has more fields populated.
+     *
+     *  Testing:
+     *
+     *      - only that specific field is updated
+     */
+    @IsTest
+    private static void testRegisterDirty_onlyDirtyField() {
+        String originalName = 'test name';
+        Opportunity opp = new Opportunity(Name = originalName, StageName = 'Open', Amount = 200, CloseDate = System.today());
+        insert opp;
+
+        opp.Name = 'Changed name';
+        opp.Amount = 300;
+
+        fflib_SObjectUnitOfWork uow = new fflib_SObjectUnitOfWork(MY_SOBJECTS);
+        uow.registerDirty( opp, new List<SObjectField> { Opportunity.Amount } );
+        uow.commitWork();
+
+        Opportunity afterUpdate = [SELECT Name, Amount FROM Opportunity WHERE Id = :opp.Id];
+        System.assertEquals( opp.Amount, afterUpdate.Amount, 'Expected Amount to be changed after registerDirty' );
+        System.assertEquals( originalName, afterUpdate.Name, 'Expected Name to be unchanged as field was not registered' );
+    }
+
     /**
      * Assert that actual events exactly match expected events (size, order and name)
      * and types match expected types