diff --git a/grader_support/gradelib.py b/grader_support/gradelib.py
index 0d05a87..aa35b3a 100644
--- a/grader_support/gradelib.py
+++ b/grader_support/gradelib.py
@@ -3,7 +3,8 @@
 import random
 import re
 import sys
-import tokenize
+from tokenize import tokenize, COMMENT, STRING
+from io import BytesIO
 
 import six
 # the run library should overwrite this with a particular random seed for the test.
@@ -238,11 +239,7 @@ def _tokens(code):
     A wrapper around tokenize.generate_tokens.
     """
     # Protect against pathological inputs: http://bugs.python.org/issue16152
-    code = code.rstrip() + "\n"
-    if isinstance(code, str):
-        code = code.encode('utf8')
-    code = "# coding: utf8\n" + code
-    toks = tokenize.generate_tokens(six.BytesIO(code).readline)
+    toks = tokenize(BytesIO(code.encode('utf-8')).readline)
     return toks
 
 def _count_tokens(code, string):
@@ -253,7 +250,7 @@ def _count_tokens(code, string):
 
     try:
         for ttyp, ttok, __, __, __ in _tokens(code):
-            if ttyp in (tokenize.COMMENT, tokenize.STRING):
+            if ttyp in (COMMENT, STRING):
                 continue
             if ttok == string:
                 count += 1
@@ -353,7 +350,7 @@ def count_non_comment_lines(at_least=None, at_most=None, exactly=None, error_msg
     def check(code):
         linenums = set()
         for ttyp, ttok, (srow, __), __, __ in _tokens(code):
-            if ttyp in (tokenize.COMMENT, tokenize.STRING):
+            if ttyp in (COMMENT, STRING):
                 # Comments and strings don't count toward line count. If a string
                 # is the only thing on a line, then it's probably a docstring, so
                 # don't count it.
@@ -530,7 +527,7 @@ def invoke_student_function(fn_name, args, environment=None, output_writer=None)
     """
     output_writer = output_writer or repr
     def doit(submission_module):
-        for name, value in environment or {}.items():
+        for name, value in (environment or {}).items():
             setattr(submission_module, name, value)
         fn = getattr(submission_module, fn_name)
         print(output_writer(fn(*args)))
diff --git a/tests/test_manager.py b/tests/test_manager.py
index 8e12518..6093bf2 100644
--- a/tests/test_manager.py
+++ b/tests/test_manager.py
@@ -67,7 +67,8 @@ def test_codejail_config(self):
                 "VMEM": 1024
             }
         }
-        self.m.enable_codejail(config)
+        codejail_return = self.m.enable_codejail(config)
+        self.assertEqual(codejail_return, config["name"])
         self.assertTrue(codejail.jail_code.is_configured("python"))
         self.m.enable_codejail({
             "name": "other-python",
diff --git a/xqueue_watcher/manager.py b/xqueue_watcher/manager.py
index e9c0da1..ad9c60a 100644
--- a/xqueue_watcher/manager.py
+++ b/xqueue_watcher/manager.py
@@ -119,8 +119,8 @@ def enable_codejail(self, codejail_config):
         user = codejail_config.get('user', getpass.getuser())
         jail_code.configure(name, bin_path, user=user)
         limits = codejail_config.get("limits", {})
-        for name, value in limits.items():
-            jail_code.set_limit(name, value)
+        for limit_name, value in limits.items():
+            jail_code.set_limit(limit_name, value)
         self.log.info("configured codejail -> %s %s %s", name, bin_path, user)
         return name