Use only where appropriate

2 files changed, 10 insertions, 11 deletions

diff --git a/src/test/java/org/zwobble/couscous/tests/frontends/java/ExpressionReadingTests.java b/src/test/java/org/zwobble/couscous/tests/frontends/java/ExpressionReadingTests.java
index e8b391e..e7b2f66 100644
--- a/src/test/java/org/zwobble/couscous/tests/frontends/java/ExpressionReadingTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/frontends/java/ExpressionReadingTests.java
@@ -33,6 +33,7 @@ import static org.zwobble.couscous.tests.frontends.java.JavaReading.*;
 import static org.zwobble.couscous.tests.frontends.java.NodeMatchers.*;
 import static org.zwobble.couscous.types.ParameterizedType.parameterizedType;
 import static org.zwobble.couscous.types.TypeParameter.typeParameter;
+import static org.zwobble.couscous.util.ExtraIterables.only;
 import static org.zwobble.couscous.util.ExtraLists.list;
 
 public class ExpressionReadingTests {
@@ -335,7 +336,7 @@ public class ExpressionReadingTests {
             "java.util.Objects.toString(42)");
         assertEquals(
             typeCoercion(literal(42), Types.OBJECT),
-            expression.getArguments().get(0));
+            only(expression.getArguments()));
     }
 
     @Test
@@ -465,7 +466,7 @@ public class ExpressionReadingTests {
                 "    return name = \"blah\";" +
                 "}");
 
-        ReturnNode returnNode = (ReturnNode) classNode.getMethods().get(0).getBody().get().get(0);
+        ReturnNode returnNode = (ReturnNode) only(only(classNode.getMethods()).getBody().get());
         assertEquals(
             assign(
                 fieldAccess(
@@ -484,7 +485,7 @@ public class ExpressionReadingTests {
                 "    return value = 4;" +
                 "}");
 
-        ReturnNode returnNode = (ReturnNode) classNode.getMethods().get(0).getBody().get().get(0);
+        ReturnNode returnNode = (ReturnNode) only(only(classNode.getMethods()).getBody().get());
         assertEquals(
             assign(
                 fieldAccess(
diff --git a/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java b/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java
index eb6e839..e2d3d06 100644
--- a/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java
+++ b/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java
@@ -7,9 +7,8 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.List;
 
-import static org.hamcrest.Matchers.hasSize;
-import static org.junit.Assert.assertThat;
 import static org.zwobble.couscous.tests.util.ExtraFiles.deleteRecursively;
+import static org.zwobble.couscous.util.ExtraIterables.only;
 import static org.zwobble.couscous.util.ExtraLists.list;
 
 public class JavaReading {
@@ -20,7 +19,8 @@ public class JavaReading {
                 "}";
 
         ClassNode classNode = readClass(javaClass);
-        ReturnNode returnStatement = (ReturnNode) classNode.getMethods().get(0).getBody().get().get(0);
+        MethodNode method = only(classNode.getMethods());
+        ReturnNode returnStatement = (ReturnNode) only(method.getBody().get());
         return returnStatement.getValue();
     }
 
@@ -38,20 +38,18 @@ public class JavaReading {
     }
 
     static StatementNode readStatement(String returnType, String statementSource) {
-        // TODO: use only()
-        return readStatements(returnType, statementSource).get(0);
+        return only(readStatements(returnType, statementSource));
     }
 
     static List<StatementNode> readStatements(String returnType, String statementsSource) {
         String javaClass = generateMethodSource(returnType, statementsSource);
         ClassNode classNode = readClass(javaClass);
-        return classNode.getMethods().get(0).getBody().get();
+        return only(classNode.getMethods()).getBody().get();
     }
 
     static ClassNode readClass(String classBody) {
         List<TypeNode> classes = readTypes(classBody);
-        assertThat(classes, hasSize(1));
-        return (ClassNode)classes.get(0);
+        return (ClassNode) only(classes);
     }
 
     static List<TypeNode> readTypes(String classBody) {