diff --git a/test/test_py_codegen.py b/test/test_py_codegen.py
index b418cc449cd929abd57119ac5bccfc12e1280d38..0015e7a719fcaed2fed754fbbbe6b1b7346c3f0d 100644
--- a/test/test_py_codegen.py
+++ b/test/test_py_codegen.py
@@ -8,14 +8,13 @@ import sys
 
 def test_pickling_with_module_import():
     cg = codegen.PythonCodeGenerator()
-    cg("import pytools.py_codegen")
+    cg("import pytools")
     cg("import math as m")
 
     import pickle
     mod = pickle.loads(pickle.dumps(cg.get_picklable_module()))
 
     assert mod.mod_globals["pytools"] is pytools
-    assert mod.mod_globals["pytools"].py_codegen is pytools.py_codegen
 
     import math
     assert mod.mod_globals["m"] is math