--- a/test/testpyy.py	Sun Dec 22 15:50:59 2013 +0100
+++ b/test/testpyy.py	Tue Dec 31 12:38:15 2013 +0100
@@ -1,10 +1,25 @@
-import unittest, pprint
-from pyyacc import Grammar, Item, ParserGenerationException, ParserException, EPS, EOF
+import unittest
+from pyyacc import Grammar, Item, ParserGenerationException, ParserException
+from pyyacc import EPS, EOF
 from ppci import Token
 
-def genTokens(lst):
-    for t in lst:
-        yield Token(t, t)
+
+class genTokens:
+    def __init__(self, lst):
+        def tokGen():
+            for t in lst:
+                yield Token(t, t)
+            while True:
+                yield Token(EOF, EOF)
+        self.tokens = tokGen()
+        self.token = self.tokens.__next__()
+
+    def next_token(self):
+        t = self.token
+        if t.typ != EOF:
+            self.token = self.tokens.__next__()
+        return t
+
 
 class testLR(unittest.TestCase):
     """ Test basic LR(1) parser generator constructs """
@@ -25,6 +40,7 @@
         p = g.genParser()
         # 4. feed input:
         p.parse(tokens)
+
     def testReduceReduceConflict(self):
         """ Check if a reduce-reduce conflict is detected """
         # Define a grammar with an obvious reduce-reduce conflict:
@@ -37,6 +53,7 @@
         g.start_symbol = 'goal'
         with self.assertRaises(ParserGenerationException):
             p = g.genParser()
+
     def testShiftReduceConflict(self):
         """ Must be handled automatically by doing shift """
         g = Grammar([EOF, 'if', 'then', 'else', 'ass'])
@@ -60,6 +77,7 @@
         g.start_symbol = 'goal'
         with self.assertRaises(ParserGenerationException):
             g.genParser()
+
     def testRedefineTerminal(self):
         """ Test correct behavior when a terminal is redefined """
         g = Grammar([EOF, 'b', 'c'])
@@ -69,6 +87,7 @@
         g.add_production('a', ['c'])
         g.start_symbol = 'goal'
         g.genParser()
+
     def testEmpty(self):
         """ Test empty token stream """
         g = Grammar([','])
@@ -78,7 +97,7 @@
         tokens = genTokens([])
         with self.assertRaises(ParserException):
             p.parse(tokens)
-        
+
     def testEps(self):
         """ Test epsilon terminal """
         g = Grammar(['a', 'b'])
@@ -109,9 +128,9 @@
         self.cb_called = False
         def cb(a, c, b):
             self.cb_called = True
-            self.assertEqual(a, 'a')
-            self.assertEqual(b, 'b')
-            self.assertEqual(c, 'c')
+            self.assertEqual(a.val, 'a')
+            self.assertEqual(b.val, 'b')
+            self.assertEqual(c.val, 'c')
         g = Grammar(['a', 'b', 'c'])
         g.add_production('goal', ['a', 'c', 'b'], cb)
         g.start_symbol = 'goal'
@@ -147,7 +166,8 @@
         # Must result in 12 sets:
         self.assertEqual(len(s), 24)
 
-class testPG(unittest.TestCase):
+
+class testParserGenerator(unittest.TestCase):
     """ Tests several parts of the parser generator """
     def setUp(self):
         g = Grammar(['(', ')'])
@@ -216,5 +236,3 @@
 
 if __name__ == '__main__':
     unittest.main()
-
-