diff --git a/java/tdd-training/src/main/java/Primes.java b/java/tdd-training/src/main/java/Primes.java
index c83f12e..c0058ff 100644
--- a/java/tdd-training/src/main/java/Primes.java
+++ b/java/tdd-training/src/main/java/Primes.java
@@ -1,8 +1,9 @@
+import java.util.Arrays;
+import java.util.List;
 
 public class Primes {
-
-
-    public static int[] of(int i) {
-        return new int[]{1};
+    
+    public static List<Integer> of(int i) {
+        return Arrays.asList(1);
     }
 }
diff --git a/java/tdd-training/src/test/java/PrimesTest.java b/java/tdd-training/src/test/java/PrimesTest.java
index d88df33..1b3e325 100644
--- a/java/tdd-training/src/test/java/PrimesTest.java
+++ b/java/tdd-training/src/test/java/PrimesTest.java
@@ -1,11 +1,12 @@
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.Arrays;
+
 public class PrimesTest {
 
     @Test
     public void prime_of_1() throws Exception {
-        int[] expected = new int[]{1};
-        Assert.assertArrayEquals(expected, Primes.of(1));
+        Assert.assertEquals(Arrays.asList(1), Primes.of(1));
     }
-}
\ No newline at end of file
+}