1 files changed, 41 insertions, 33 deletions

diff --git a/chapter01/src/test/java/com/seaofnodes/simple/Chapter01Test.java b/chapter01/src/test/java/com/seaofnodes/simple/Chapter01Test.java
index 14c22437..ce8c2822 100644
--- a/chapter01/src/test/java/com/seaofnodes/simple/Chapter01Test.java
+++ b/chapter01/src/test/java/com/seaofnodes/simple/Chapter01Test.java
@@ -1,39 +1,42 @@
 package com.seaofnodes.simple;
 
-import com.seaofnodes.simple.node.*;
+import com.seaofnodes.simple.node.ConstantNode;
+import com.seaofnodes.simple.node.Node;
+import com.seaofnodes.simple.node.ReturnNode;
+import com.seaofnodes.simple.node.StartNode;
+import org.junit.Rule;
 import org.junit.Test;
-import static org.junit.Assert.*;
+import org.junit.rules.ExpectedException;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
 
 public class Chapter01Test {
 
+    @Rule
+    public ExpectedException exceptionRule = ExpectedException.none();
+
     @Test
     public void testSimpleProgram() {
-        Parser parser = new Parser();
-        StartNode startNode = parser.parse("return 1;");
-        assertNotNull(startNode);
-        assertEquals(2, startNode.nOuts());
-        for (int i = 0; i < startNode.nOuts(); i++) {
-            switch( startNode.out(i) ) {
-            case ReturnNode ret:
-                assertEquals(2, ret.nIns());
-                assertEquals(startNode, ret.in(0));
-                assertTrue(ret.in(1) instanceof ConstantNode);
-                break;
-            case ConstantNode con:
-                assertEquals(1, con._value);
-                assertEquals(1, con.nIns());
-                assertEquals(startNode, con.in(0));
-                break;
-            default:
-                throw new IllegalStateException();
-            }
+        Parser parser = new Parser("return 1;");
+        ReturnNode ret = parser.parse();
+        StartNode start = Parser.START;
+
+        assertEquals(start, ret.ctrl());
+        Node expr = ret.expr();
+        if (expr instanceof ConstantNode con) {
+            assertEquals(start, con.in(0));
+            assertEquals(1, con._value);
+        } else {
+            fail();
         }
     }
 
     @Test
     public void testZero() {
-        Parser parser = new Parser();
-        StartNode start = parser.parse("return 0;");
+        Parser parser = new Parser("return 0;");
+        parser.parse();
+        StartNode start = Parser.START;
         for( Node use : start._outputs )
             if( use instanceof ConstantNode con )
                 assertEquals(0,con._value);
@@ -41,32 +44,37 @@ public class Chapter01Test {
 
     @Test
     public void testBad1() {
-        try { new Parser().parse("ret"); }
-        catch( RuntimeException e ) { assertEquals("syntax error, expected return: ret",e.getMessage()); }
+        exceptionRule.expect(RuntimeException.class);
+        exceptionRule.expectMessage("Syntax error, expected return: ret");
+        new Parser("ret").parse();
     }
 
     @Test
     public void testBad2() {
-        try { new Parser().parse("return 0123;"); }
-        catch( RuntimeException e ) { assertEquals("syntax error: integer values cannot start with '0'",e.getMessage()); }
+        exceptionRule.expect(RuntimeException.class);
+        exceptionRule.expectMessage("Syntax error: integer values cannot start with '0'");
+        new Parser("return 0123;").parse();
     }
 
     @Test
     public void testBad3() {
-        try { new Parser().parse("return --12;"); }
-        catch( RuntimeException e ) { assertEquals("syntax error, expected integer literal: -",e.getMessage()); }
+        exceptionRule.expect(RuntimeException.class);
+        exceptionRule.expectMessage("Syntax error, expected integer literal");
+        new Parser("return --12;").parse();
     }
 
     @Test
     public void testBad4() {
-        try { new Parser().parse("return 100"); }
-        catch( RuntimeException e ) { assertTrue(e.getMessage().contains("syntax error, expected ; got :")); }
+        exceptionRule.expect(RuntimeException.class);
+        exceptionRule.expectMessage("Syntax error, expected ;:");
+        new Parser("return 100").parse();
     }
 
     // Negative numbers require unary operator support that is not in scope
     @Test
     public void testBad5() {
-        try { new Parser().parse("return -100;"); }
-        catch( RuntimeException e ) { assertEquals("syntax error, expected integer literal: -", e.getMessage()); }
+        exceptionRule.expect(RuntimeException.class);
+        exceptionRule.expectMessage("Syntax error, expected integer literal");
+        new Parser("return -100;").parse();
     }
 }