Update test to current lexer functions

1 files changed, 4 insertions, 9 deletions

diff --git a/tests/test_tokenize.py b/tests/test_tokenize.py
index 4bcfd49..3580bc8 100644
--- a/tests/test_tokenize.py
+++ b/tests/test_tokenize.py
@@ -126,27 +126,22 @@ class TestTokenList(unittest.TestCase):
 class TestStream(unittest.TestCase):
     def test_simple(self):
         stream = StringIO("SELECT 1; SELECT 2;")
-        lex = lexer.Lexer()
 
-        tokens = lex.get_tokens(stream)
+        tokens = lexer.tokenize(stream)
         self.assertEqual(len(list(tokens)), 9)
 
         stream.seek(0)
-        lex.bufsize = 4
-        tokens = list(lex.get_tokens(stream))
+        tokens = list(lexer.tokenize(stream))
         self.assertEqual(len(tokens), 9)
 
         stream.seek(0)
-        lex.bufsize = len(stream.getvalue())
-        tokens = list(lex.get_tokens(stream))
+        tokens = list(lexer.tokenize(stream))
         self.assertEqual(len(tokens), 9)
 
     def test_error(self):
         stream = StringIO("FOOBAR{")
 
-        lex = lexer.Lexer()
-        lex.bufsize = 4
-        tokens = list(lex.get_tokens(stream))
+        tokens = list(lexer.tokenize(stream))
         self.assertEqual(len(tokens), 2)
         self.assertEqual(tokens[1][0], T.Error)