Validate Rail Fence cipher rail counts

Author: krotname

src/main/java/kyu3/RailFenceCipher.java

@@ -11,6 +11,7 @@ public class RailFenceCipher {
     //3 https://www.codewars.com/kata/58c5577d61aefcf3ff000081/train/java
 
     static String encode(String s, int n) {
+        validateRailCount(n);
         // Distributes characters across rails using a periodic index that walks
         // up and down between boundary rails, then concatenates each rail content.
         Map<Integer, StringBuilder> map = new TreeMap<>();
@@ -37,6 +38,7 @@ static String encode(String s, int n) {
     }
 
     static String decode(String s, int n) {
+        validateRailCount(n);
         // First computes the exact size of each rail, splits ciphertext into
         // contiguous rail segments, then reconstructs plaintext by replaying the
         // same rail traversal cycle.
@@ -80,6 +82,12 @@ static String decode(String s, int n) {
         return result.toString();
     }
 
+    private static void validateRailCount(int railCount) {
+        if (railCount < 2) {
+            throw new IllegalArgumentException("Rail count must be at least 2");
+        }
+    }
+
     private static int sumArr(int[] arr, int i) {
         // Sum widths of all previous rails to determine the start offset
         // of the current rail segment.

src/test/java/kyu3/RailFenceCipherTest.java

@@ -1,6 +1,7 @@
 package kyu3;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Tag;
@@ -28,11 +29,28 @@ void shouldRoundTripPlainTextForDifferentRails(String text, int rails) {
         assertEquals(text, RailFenceCipher.decode(encoded, rails));
     }
 
+    @ParameterizedTest
+    @MethodSource("invalidRailCounts")
+    void shouldRejectInvalidRailCounts(int rails) {
+        assertThrows(IllegalArgumentException.class,
+                () -> RailFenceCipher.encode("ABC", rails));
+        assertThrows(IllegalArgumentException.class,
+                () -> RailFenceCipher.decode("ABC", rails));
+    }
+
     private static Stream<Arguments> roundTripCases() {
         return Stream.of(
                 Arguments.of("", 2),
                 Arguments.of("Hello, World!", 4),
                 Arguments.of("Rail fence cipher keeps punctuation.", 5)
         );
     }
+
+    private static Stream<Arguments> invalidRailCounts() {
+        return Stream.of(
+                Arguments.of(-1),
+                Arguments.of(0),
+                Arguments.of(1)
+        );
+    }
 }