Move zones() up to simplify

2 files changed, 20 insertions, 27 deletions

diff --git a/config-model-api/src/main/java/com/yahoo/config/application/api/DeploymentSpec.java b/config-model-api/src/main/java/com/yahoo/config/application/api/DeploymentSpec.java
index fc7344aa9e8..8530679e82f 100644
--- a/config-model-api/src/main/java/com/yahoo/config/application/api/DeploymentSpec.java
+++ b/config-model-api/src/main/java/com/yahoo/config/application/api/DeploymentSpec.java
@@ -15,6 +15,7 @@ import java.io.Reader;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -76,13 +77,9 @@ public class DeploymentSpec {
     private void validateZones(List<Step> steps) {
         Set<DeclaredZone> zones = new HashSet<>();
 
-        steps.stream().filter(step -> step instanceof DeclaredZone)
-                      .map(DeclaredZone.class::cast)
-                      .forEach(zone -> ensureUnique(zone, zones));
-        steps.stream().filter(step -> step instanceof ParallelZones)
-                      .map(ParallelZones.class::cast)
-                      .flatMap(parallelZones -> parallelZones.zones().stream())
-                      .forEach(zone -> ensureUnique(zone, zones));
+        for (Step step : steps)
+            for (DeclaredZone zone : step.zones())
+                ensureUnique(zone, zones);
     }
     
     private void ensureUnique(DeclaredZone zone, Set<DeclaredZone> zones) {
@@ -320,6 +317,9 @@ public class DeploymentSpec {
         /** Returns whether this step deploys to the given environment, and (if specified) region */
         public abstract boolean deploysTo(Environment environment, Optional<RegionName> region);
 
+        /** Returns the zones deployed to in this step */
+        public List<DeclaredZone> zones() { return Collections.emptyList(); }
+
     }
 
     /** A deployment step which is to wait for some time before progressing to the next step */
@@ -370,6 +370,9 @@ public class DeploymentSpec {
         public boolean active() { return active; }
 
         @Override
+        public List<DeclaredZone> zones() { return Collections.singletonList(this); }
+
+        @Override
         public boolean deploysTo(Environment environment, Optional<RegionName> region) {
             if (environment != this.environment) return false;
             if (region.isPresent() && ! region.equals(this.region)) return false;
@@ -407,7 +410,7 @@ public class DeploymentSpec {
             this.zones = ImmutableList.copyOf(zones);
         }
 
-        /** The list of zones to deploy in */
+        @Override
         public List<DeclaredZone> zones() { return this.zones; }
 
         @Override
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentOrder.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentOrder.java
index afa9c219048..45979b597f0 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentOrder.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentOrder.java
@@ -47,7 +47,7 @@ public class DeploymentOrder {
         // At this point we have deployed to system test, so deployment spec is available
         List<DeploymentSpec.Step> deploymentSteps = deploymentSteps(application);
         Optional<DeploymentSpec.Step> currentStep = fromJob(job, application);
-        if (!currentStep.isPresent()) {
+        if ( ! currentStep.isPresent()) {
             return Collections.emptyList();
         }
 
@@ -65,15 +65,9 @@ public class DeploymentOrder {
         }
 
         DeploymentSpec.Step nextStep = deploymentSteps.get(currentIndex + 1);
-        if (nextStep instanceof DeploymentSpec.DeclaredZone) {
-            return Collections.singletonList(toJob((DeploymentSpec.DeclaredZone) nextStep));
-        } else if (nextStep instanceof DeploymentSpec.ParallelZones) {
-            return ((DeploymentSpec.ParallelZones) nextStep).zones().stream()
-                    .map(this::toJob)
-                    .collect(collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
-        } else {
-            throw new IllegalStateException("Unexpected step type: " + nextStep.getClass());
-        }
+        return nextStep.zones().stream()
+                .map(this::toJob)
+                .collect(collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
     }
 
     /** Returns whether the given job is first in a deployment */
@@ -96,15 +90,11 @@ public class DeploymentOrder {
         if (deploymentSpec.steps().isEmpty()) {
             return Arrays.asList(JobType.systemTest, JobType.stagingTest);
         }
-        List<JobType> jobs = new ArrayList<>();
-        for (DeploymentSpec.Step step : deploymentSpec.steps()) {
-            if (step instanceof DeploymentSpec.DeclaredZone) {
-                jobs.add(toJob((DeploymentSpec.DeclaredZone) step));
-            } else if (step instanceof DeploymentSpec.ParallelZones) {
-                ((DeploymentSpec.ParallelZones) step).zones().forEach(zone -> jobs.add(toJob(zone)));
-            }
-        }
-        return Collections.unmodifiableList(jobs);
+
+        return deploymentSpec.steps().stream()
+                .flatMap(step -> step.zones().stream())
+                .map(this::toJob)
+                .collect(Collectors.collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
     }
 
     /** Resolve deployment step from job */