package water.runner;

import org.junit.Ignore;
import water.*;

import java.util.Set;

@Ignore
public class CleanNewKeysTask extends KeysMRTask<CleanNewKeysTask> {

    @Override
    protected void setupLocal() {
        final Set<Key> initKeys = LocalTestRuntime.initKeys;
        final Set<Key> actualKeys = H2O.localKeySet();
        for (Key actualKey : actualKeys){

            final Value value = Value.STORE_get(actualKey);
            if (initKeys.contains(actualKey) || isIgnorableKeyLeak(actualKey, value)) continue;
            if (!(value.get() instanceof Keyed)) {
                // Keyed objects might override remove_impl to excerscise their own removal strategy
                // Non-keyed objects should just be removed from the DKV
                DKV.remove(actualKey);
            } else {
                actualKey.remove();
            }
        }
        
    }

}

@@ -219,668 +219,73 @@
     
-final Set%3CKey%3E initKeys = LocalTestRuntime.initKeys;%0A        final Set%3CKey%3E actualKeys = H2O.localKeySet();%0A        for (Key actualKey : actualKeys)%7B%0A%0A            final Value value = Value.STORE_get(actualKey);%0A            if (initKeys.contains(actualKey) %7C%7C isIgnorableKeyLeak(actualKey, value)) continue;%0A            if (!(value.get() instanceof Keyed)) %7B%0A                // Keyed objects might override remove_impl to excerscise their own removal strategy%0A                // Non-keyed objects should just be removed from the DKV%0A                DKV.remove(actualKey);%0A            %7D else %7B%0A                actualKey.remove();%0A            %7D%0A        %7D%0A        
+DKVManager.retain(LocalTestRuntime.initKeys.toArray(new Key%5B0%5D));
 %0A   

Use DKVManager to clean newly added keys.