diff --git a/src/main/java/kyu3/RailFenceCipher.java b/src/main/java/kyu3/RailFenceCipher.java index 4c5c22f..c7008a5 100644 --- a/src/main/java/kyu3/RailFenceCipher.java +++ b/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 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. diff --git a/src/test/java/kyu3/RailFenceCipherTest.java b/src/test/java/kyu3/RailFenceCipherTest.java index ccc1fcc..4693629 100644 --- a/src/test/java/kyu3/RailFenceCipherTest.java +++ b/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,6 +29,15 @@ 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 roundTripCases() { return Stream.of( Arguments.of("", 2), @@ -35,4 +45,12 @@ private static Stream roundTripCases() { Arguments.of("Rail fence cipher keeps punctuation.", 5) ); } + + private static Stream invalidRailCounts() { + return Stream.of( + Arguments.of(-1), + Arguments.of(0), + Arguments.of(1) + ); + } }