--- a/test/testasm.py	Sat Apr 26 17:41:56 2014 +0200
+++ b/test/testasm.py	Sun Apr 27 12:24:21 2014 +0200
@@ -2,7 +2,7 @@
 
 import unittest
 from ppci import CompilerError
-from ppci.assembler import tokenize
+from ppci.assembler import AsmLexer
 from ppci.objectfile import ObjectFile
 from ppci.outstream import BinaryOutputStream
 from ppci.target.basetarget import Label
@@ -12,26 +12,38 @@
 class AssemblerLexingCase(unittest.TestCase):
     """ Tests the assemblers lexer """
 
+    def setUp(self):
+        self.lexer = AsmLexer([])
+
+    def do(self, asmline, toks):
+        output = []
+        self.lexer.feed(asmline)
+        while 'EOF' not in output:
+            output.append(self.lexer.next_token().typ)
+        self.assertSequenceEqual(toks, output)
+
     def testLex0(self):
         """ Check if the lexer is OK """
-        asmline, toks = 'mov rax, rbx ', ['ID', 'ID', ',', 'ID', 'EOF']
-        self.assertSequenceEqual([tok.typ for tok in tokenize(asmline, [])], toks)
+        asmline = 'mov rax, rbx '
+        toks = ['ID', 'ID', ',', 'ID', 'EOF']
+        self.do(asmline, toks)
 
     def testLex1(self):
         """ Test if lexer correctly maps some tokens """
-        asmline, toks = 'lab1: mov rax, rbx ', ['ID', ':', 'ID', 'ID', ',', 'ID', 'EOF']
-        self.assertSequenceEqual([tok.typ for tok in tokenize(asmline, [])], toks)
+        asmline = 'lab1: mov rax, rbx '
+        toks = ['ID', ':', 'ID', 'ID', ',', 'ID', 'EOF']
+        self.do(asmline, toks)
 
     def testLex2(self):
         """ Test if lexer correctly maps some tokens """
         asmline, toks = 'mov 3.13 0xC 13', ['ID', 'REAL', 'val5', 'val5', 'EOF']
-        self.assertSequenceEqual([tok.typ for tok in tokenize(asmline, [])], toks)
+        self.do(asmline, toks)
 
     def testLex3(self):
         """ Test if lexer fails on a token that is invalid """
         asmline = '0z4: mov rax, rbx $ '
         with self.assertRaises(CompilerError):
-            list(tokenize(asmline, []))
+            self.do(asmline, [])
 
 
 class OustreamTestCase(unittest.TestCase):