Simplify

2 files changed, 20 insertions, 43 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/flag/Flag.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/flag/Flag.java
index 635b1b012da..a22e1dea024 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/flag/Flag.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/flag/Flag.java
@@ -62,41 +62,30 @@ public class Flag {
     }
 
     /** Returns a copy of this with this flag enabled for all dimensions */
-    public Flag enable() {
-        return new Flag(id, true, hostnames, applications);
-    }
-
-    /** Returns a copy of this with flag enabled for given hostname */
-    public Flag enable(HostName hostname) {
-        Set<String> hostnames = new LinkedHashSet<>(this.hostnames);
-        hostnames.add(hostname.value());
+    public Flag withEnabled(boolean enabled) {
         return new Flag(id, enabled, hostnames, applications);
     }
 
-    /** Returns a copy of this with flag enabled for given application */
-    public Flag enable(ApplicationId application) {
-        Set<ApplicationId> applications = new LinkedHashSet<>(this.applications);
-        applications.add(application);
-        return new Flag(id, enabled, hostnames, applications);
-    }
-
-    /** Returns a copy of this with flag disabled for given hostname */
-    public Flag disable(HostName hostname) {
+    /** Returns a copy of this with enabled set for hostname */
+    public Flag withEnabled(HostName hostname, boolean enabled) {
         Set<String> hostnames = new LinkedHashSet<>(this.hostnames);
-        hostnames.remove(hostname.value());
-        return new Flag(id, enabled, hostnames, applications);
+        if (enabled) {
+            hostnames.add(hostname.value());
+        } else {
+            hostnames.remove(hostname.value());
+        }
+        return new Flag(id, this.enabled, hostnames, applications);
     }
 
-    /** Returns a copy of this with flag disabled for given application */
-    public Flag disable(ApplicationId application) {
+    /** Returns a copy of this with enabled set for application */
+    public Flag withEnabled(ApplicationId application, boolean enabled) {
         Set<ApplicationId> applications = new LinkedHashSet<>(this.applications);
-        applications.remove(application);
-        return new Flag(id, enabled, hostnames, applications);
-    }
-
-    /** Returns a copy of this with this flag disabled in all dimensions */
-    public Flag disable() {
-        return new Flag(id, false, hostnames, applications);
+        if (enabled) {
+            applications.add(application);
+        } else {
+            applications.remove(application);
+        }
+        return new Flag(id, this.enabled, hostnames, applications);
     }
 
     @Override
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/flag/Flags.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/flag/Flags.java
index 31344d24c9e..b4ecf415ede 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/flag/Flags.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/flag/Flags.java
@@ -48,29 +48,17 @@ public class Flags {
 
     /** Enable feature flag in this node repository */
     public void setEnabled(FlagId flag, boolean enabled) {
-        if (enabled) {
-            write(flag, Flag::enable);
-        } else {
-            write(flag, Flag::disable);
-        }
+        write(flag, (f) -> f.withEnabled(enabled));
     }
 
     /** Enable feature flag for given application */
     public void setEnabled(FlagId flag, ApplicationId application, boolean enabled) {
-        if (enabled) {
-            write(flag, (f) -> f.enable(application));
-        } else {
-            write(flag, (f) -> f.disable(application));
-        }
+        write(flag, (f) -> f.withEnabled(application, enabled));
     }
 
     /** Enable feature flag for given node */
     public void setEnabled(FlagId flag, HostName hostname, boolean enabled) {
-        if (enabled) {
-            write(flag, (f) -> f.enable(hostname));
-        } else {
-            write(flag, (f) -> f.disable(hostname));
-        }
+        write(flag, (f) -> f.withEnabled(hostname, enabled));
     }
 
     private void write(FlagId id, Function<Flag, Flag> updateFunc) {