1 files changed, 17 insertions, 8 deletions
diff --git a/src/main/java/design/model/Play.java b/src/main/java/design/model/Play.java
index ca6dac2..2e083b9 100644
--- a/src/main/java/design/model/Play.java
+++ b/src/main/java/design/model/Play.java
@@ -6,22 +6,22 @@ import java.util.ArrayList;
 import java.util.List;
 
 public class Play {
-    private final int distance;
+    private int holeNumber;
     private final List<Swing> swings;
 
     @JsonCreator
-    private Play(int distance, List<Swing> swings) {
-        this.distance = distance;
-        this.swings = swings;
+    public Play(int holeNumber, List<Swing> swings) {
+        this.holeNumber = holeNumber;
+        this.swings = (swings != null) ? swings : new ArrayList<>();
     }
 
-    public Play(int distance) {
+    public Play(int holeNumber) {
+        this.holeNumber = holeNumber;
         this.swings = new ArrayList<>();
-        this.distance = distance;
     }
 
-    public int getDistance() {
-        return distance;
+    public int getHoleNumber() {
+        return holeNumber;
     }
 
     public void addSwing(Swing swing) {
@@ -35,4 +35,13 @@ public class Play {
     public int getSwingCount() {
         return swings.size();
     }
+
+    // Returns the sum of all swing distances
+    public int getDistance() {
+        int total = 0;
+        for (Swing s : swings) {
+            total += s.getDistance();
+        }
+        return total;
+    }
 }