diff --git a/pystencils_tests/test_address_of.py b/pystencils_tests/test_address_of.py
index 8de48e2bb0783dc66452a0bf31a7a7529d8d8292..5b297fc63a8e088d858d2abd939b289961c2d7f6 100644
--- a/pystencils_tests/test_address_of.py
+++ b/pystencils_tests/test_address_of.py
@@ -3,19 +3,18 @@
 Test of pystencils.data_types.address_of
 """
 
-from pystencils.data_types import address_of, cast_func, PointerType
 import pystencils
+from pystencils.data_types import PointerType, address_of, cast_func, create_type
 from pystencils.simp.simplifications import sympy_cse
-import sympy
 
 
 def test_address_of():
     x, y = pystencils.fields('x,y: int64[2d]')
-    s = pystencils.TypedSymbol('s', PointerType('int64'))
+    s = pystencils.TypedSymbol('s', PointerType(create_type('int64')))
 
     assignments = pystencils.AssignmentCollection({
         s: address_of(x[0, 0]),
-        y[0, 0]: cast_func(s, 'int64')
+        y[0, 0]: cast_func(s, create_type('int64'))
     }, {})
 
     ast = pystencils.create_kernel(assignments)
@@ -23,7 +22,7 @@ def test_address_of():
     print(code)
 
     assignments = pystencils.AssignmentCollection({
-        y[0, 0]: cast_func(address_of(x[0, 0]), 'int64')
+        y[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64'))
     }, {})
 
     ast = pystencils.create_kernel(assignments)
@@ -33,11 +32,11 @@ def test_address_of():
 
 def test_address_of_with_cse():
     x, y = pystencils.fields('x,y: int64[2d]')
-    s = pystencils.TypedSymbol('s', PointerType('int64'))
+    s = pystencils.TypedSymbol('s', PointerType(create_type('int64')))
 
     assignments = pystencils.AssignmentCollection({
-        y[0, 0]: cast_func(address_of(x[0, 0]), 'int64'),
-        x[0, 0]: cast_func(address_of(x[0, 0]), 'int64') + 1
+        y[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64')) + s,
+        x[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64')) + 1
     }, {})
 
     ast = pystencils.create_kernel(assignments)