diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 7a7309e..b0ef27d 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -14,13 +14,20 @@ defaults:
 
 jobs:
   linux-unittests:
-    name: Unit tests Linux - Python ${{ matrix.PYTHON_VERSION }}
+    name: Unit tests Linux - ${{ matrix.PYTHON_VERSION }} ${{ matrix.POLARS_VERSION }}
     timeout-minutes: 15
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        PYTHON_VERSION: ['3.9', '3.10', '3.11']
+        include:
+          - { PYTHON_VERSION: 'python=3.9', POLARS_VERSION: 'polars=0.14.28' }
+          - { PYTHON_VERSION: 'python=3.9', POLARS_VERSION: 'polars=0.15' }
+          - { PYTHON_VERSION: 'python=3.9', POLARS_VERSION: 'polars=0.16' }
+          - { PYTHON_VERSION: 'python=3.9', POLARS_VERSION: 'polars=0.17' }
+          - { PYTHON_VERSION: 'python=3.9', POLARS_VERSION: 'polars=0.18' }
+          - { PYTHON_VERSION: 'python=3.10', POLARS_VERSION: '' }
+          - { PYTHON_VERSION: 'python=3.11', POLARS_VERSION: '' }
     steps:
       - uses: actions/checkout@v3
       # TODO: move to action once it is available
@@ -28,17 +35,17 @@ jobs:
         run: |
           curl -fsSL https://raw.githubusercontent.com/prefix-dev/pixi/main/install/install.sh | bash
       - name: Install dependencies
-        # TODO: make prettier once pixi supports it
+        # TODO: make prettier once there are feature flags
         # https://github.com/prefix-dev/pixi/issues/239
         run: |
-          pixi add python=${{ matrix.PYTHON_VERSION }}
+          pixi add ${{ matrix.PYTHON_VERSION }} ${{ matrix.POLARS_VERSION }}
           pixi install
           pixi run postinstall
       - name: Run unittests
         uses: pavelzw/pytest-action@v2
         with:
           custom-pytest: pixi run pytest
-          report-title: Unit tests Linux - Python ${{ matrix.PYTHON_VERSION }}
+          report-title: Unit tests Linux - ${{ matrix.PYTHON_VERSION }} ${{ matrix.POLARS_VERSION }}
 
   pre-commit-checks:
     # TODO: switch to pixi once there is a good way
diff --git a/pixi.toml b/pixi.toml
index 746603c..fd80b10 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -9,20 +9,20 @@ channels = ["conda-forge"]
 platforms = ["linux-64", "osx-arm64", "osx-64", "win-64"]
 
 [tasks]
-"postinstall" = "pip install --no-build-isolation --no-deps --disable-pip-version-check -e ."
-"test" = "pytest"
-"lint" = "pre-commit run --all"
+postinstall = "pip install --no-build-isolation --no-deps --disable-pip-version-check -e ."
+test = "pytest"
+lint = "pre-commit run --all"
 
 [dependencies]
-python = ">= 3.9"
-"pip" = "*"
-"polars" = "0.18.8"
+python = ">=3.9"
+pip = "*"
+polars = ">=0.14.24,<0.19"
 # build
-"hatchling" = "*"
+hatchling = "*"
 # test
-"pytest" = "*"
-"pytest-md" = "*"
-"pytest-emoji" = "*"
-"hypothesis" = "*"
+pytest = "*"
+pytest-md = "*"
+pytest-emoji = "*"
+hypothesis = "*"
 # linting
-"pre-commit" = "*"
+pre-commit = "*"
diff --git a/polarify/__init__.py b/polarify/__init__.py
index bf7333f..ab72542 100644
--- a/polarify/__init__.py
+++ b/polarify/__init__.py
@@ -12,7 +12,12 @@ def transform_func_to_new_source(func) -> str:
     expr = parse_body(func_def.body)
 
     # Replace the body of the function with the parsed expr
-    func_def.body = [ast.Return(expr)]
+    # Also import polars as pl since this is used in the generated code
+    # We don't want to rely on the user having imported polars as pl
+    func_def.body = [
+        ast.Import(names=[ast.alias(name="polars", asname="pl")]),
+        ast.Return(value=expr),
+    ]
     # TODO: make this prettier
     func_def.decorator_list = []
     func_def.name += "_polarified"
diff --git a/pyproject.toml b/pyproject.toml
index c1d4d25..4e5b511 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,7 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
 ]
 dependencies = [
-    "polars == 0.18.8",
+    "polars >=0.14.24,<0.19",
 ]
 
 [project.urls]
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/functions.py b/tests/functions.py
new file mode 100644
index 0000000..7f47eca
--- /dev/null
+++ b/tests/functions.py
@@ -0,0 +1,179 @@
+# ruff: noqa
+# ruff must not change the AST of the test functions, even if they are semantically equivalent.
+
+
+def signum(x):
+    s = 0
+    if x > 0:
+        s = 1
+    elif x < 0:
+        s = -1
+    return s
+
+
+def signum_no_default(x):
+    if x > 0:
+        return 1
+    elif x < 0:
+        return -1
+    return 0
+
+
+def early_return(x):
+    if x > 0:
+        return 1
+    return 0
+
+
+def assign_both_branches(x):
+    if x > 0:
+        s = 1
+    else:
+        s = -1
+    return s
+
+
+def unary_expr(x):
+    s = -x
+    return s
+
+
+def call_target_identity(x):
+    return x
+
+
+def call_expr(x):
+    k = x * 2
+    s = call_target_identity(k + 3)
+    return s
+
+
+def if_expr(x):
+    s = 1 if x > 0 else -1
+    return s
+
+
+def if_expr2(x):
+    s = 1 + (x if x > 0 else -1)
+    return s
+
+
+def if_expr3(x):
+    s = 1 + ((3 if x < 10 else 5) if x > 0 else -1)
+    return s
+
+
+def compare_expr(x):
+    if (0 < x) & (x < 10):
+        s = 1
+    else:
+        s = 2
+    return s
+
+
+def chained_compare_expr(x):
+    if 0 < x < 10:
+        s = 1
+    else:
+        s = 2
+    return s
+
+
+def walrus_expr(x):
+    if (y := x + 1) > 0:
+        s = 1
+    else:
+        s = -1
+    return s * y
+
+
+def multiple_if_else(x):
+    if x > 0:
+        s = 1
+    elif x < 0:
+        s = -1
+    else:
+        s = 0
+    return s
+
+
+def nested_if_else(x):
+    if x > 0:
+        if x > 1:
+            s = 2
+        else:
+            s = 1
+    elif x < 0:
+        s = -1
+    else:
+        s = 0
+    return s
+
+
+def nested_if_else_expr(x):
+    if x > 0:
+        s = 2 if x > 1 else 1
+    elif x < 0:
+        s = -1
+    else:
+        s = 0
+    return s
+
+
+def assignments_inside_branch(x):
+    if x > 0:
+        s = 1
+        s = s + 1
+        s = x * s
+    elif x < 0:
+        s = -1
+        s = s - 1
+        s = x
+    else:
+        s = 0
+    return s
+
+
+def override_default(x):
+    s = 0
+    if x > 0:
+        s = 10
+    return x * s
+
+
+def no_if_else(x):
+    s = x * 10
+    k = x - 3
+    k = k * 2
+    return s * k
+
+
+def two_if_expr(x):
+    a = 1 if x > 0 else 5
+    b = 2 if x < 0 else 2
+    return a + b
+
+
+functions = [
+    signum,
+    early_return,
+    assign_both_branches,
+    unary_expr,
+    call_expr,
+    if_expr,
+    if_expr2,
+    if_expr3,
+    compare_expr,
+    multiple_if_else,
+    nested_if_else,
+    nested_if_else_expr,
+    assignments_inside_branch,
+    override_default,
+    no_if_else,
+    two_if_expr,
+]
+
+xfail_functions = [
+    walrus_expr,
+    signum_no_default,
+]
diff --git a/tests/test_error_handling.py b/tests/test_error_handling.py
new file mode 100644
index 0000000..ec028c0
--- /dev/null
+++ b/tests/test_error_handling.py
@@ -0,0 +1,10 @@
+import pytest
+
+from polarify import polarify
+
+from .functions import chained_compare_expr
+
+
+def test_chained_compare_fail():
+    with pytest.raises(ValueError):
+        polarify(chained_compare_expr)
diff --git a/tests/test_parse_body.py b/tests/test_parse_body.py
index 0bdbf48..0830163 100644
--- a/tests/test_parse_body.py
+++ b/tests/test_parse_body.py
@@ -1,196 +1,19 @@
-# ruff: noqa
-
 import inspect
-import polars as pl
+
+import polars
 import pytest
 from hypothesis import given
+from hypothesis.strategies import integers
+from packaging.version import Version
+from polars import __version__ as _pl_version
 from polars.testing import assert_frame_equal
 from polars.testing.parametric import column, dataframes
-from hypothesis.strategies import integers
 
 from polarify import polarify, transform_func_to_new_source
 
+from .functions import functions, xfail_functions
 
-def signum(x):
-    s = 0
-    if x > 0:
-        s = 1
-    elif x < 0:
-        s = -1
-    return s
-
-
-def signum_no_default(x):
-    if x > 0:
-        return 1
-    elif x < 0:
-        return -1
-    return 0
-
-
-def early_return(x):
-    if x > 0:
-        return 1
-    return 0
-
-
-def assign_both_branches(x):
-    if x > 0:
-        s = 1
-    else:
-        s = -1
-    return s
-
-
-def unary_expr(x):
-    s = -x
-    return s
-
-
-def call_target_identity(x):
-    return x
-
-
-def call_expr(x):
-    k = x * 2
-    s = call_target_identity(k + 3)
-    return s
-
-
-def if_expr(x):
-    s = 1 if x > 0 else -1
-    return s
-
-
-def if_expr2(x):
-    s = 1 + (x if x > 0 else -1)
-    return s
-
-
-def if_expr3(x):
-    s = 1 + ((3 if x < 10 else 5) if x > 0 else -1)
-    return s
-
-
-def compare_expr(x):
-    if (0 < x) & (x < 10):
-        s = 1
-    else:
-        s = 2
-    return s
-
-
-def chained_compare_expr(x):
-    if 0 < x < 10:
-        s = 1
-    else:
-        s = 2
-    return s
-
-
-def test_chained_compare_fail():
-    with pytest.raises(ValueError):
-        polarify(chained_compare_expr)
-
-
-def walrus_expr(x):
-    if (y := x + 1) > 0:
-        s = 1
-    else:
-        s = -1
-    return s * y
-
-
-def multiple_if_else(x):
-    if x > 0:
-        s = 1
-    elif x < 0:
-        s = -1
-    else:
-        s = 0
-    return s
-
-
-def nested_if_else(x):
-    if x > 0:
-        if x > 1:
-            s = 2
-        else:
-            s = 1
-    elif x < 0:
-        s = -1
-    else:
-        s = 0
-    return s
-
-
-def nested_if_else_expr(x):
-    if x > 0:
-        s = 2 if x > 1 else 1
-    elif x < 0:
-        s = -1
-    else:
-        s = 0
-    return s
-
-
-def assignments_inside_branch(x):
-    if x > 0:
-        s = 1
-        s = s + 1
-        s = x * s
-    elif x < 0:
-        s = -1
-        s = s - 1
-        s = x
-    else:
-        s = 0
-    return s
-
-
-def override_default(x):
-    s = 0
-    if x > 0:
-        s = 10
-    return x * s
-
-
-def no_if_else(x):
-    s = x * 10
-    k = x - 3
-    k = k * 2
-    return s * k
-
-
-def two_if_expr(x):
-    a = 1 if x > 0 else 5
-    b = 2 if x < 0 else 2
-    return a + b
-
-
-functions = [
-    signum,
-    early_return,
-    assign_both_branches,
-    unary_expr,
-    call_expr,
-    if_expr,
-    if_expr2,
-    if_expr3,
-    compare_expr,
-    multiple_if_else,
-    nested_if_else,
-    nested_if_else_expr,
-    assignments_inside_branch,
-    override_default,
-    no_if_else,
-    two_if_expr,
-]
-
-xfail_functions = [
-    walrus_expr,
-    signum_no_default,
-]
+pl_version = Version(_pl_version)
 
 
 @pytest.fixture(
@@ -208,19 +31,24 @@ def test_funcs(request):
     # build ast from transformed function as format as string
     transformed_func_unparsed = transform_func_to_new_source(original_func)
     print(
-        (
-            f"Original:\n{original_func_unparsed}\n"
-            f"Transformed:\n{transformed_func_unparsed}"
-        )
+        f"Original:\n{original_func_unparsed}\n"
+        f"Transformed:\n{transformed_func_unparsed}"
     )
     return transformed_func, original_func
 
 
+# chunking + apply is broken for polars < 0.18.1
+# https://github.com/pola-rs/polars/pull/9211
+# only relevant for our test setup, not for the library itself
 @given(
-    df=dataframes(column("x", dtype=pl.Int64, strategy=integers(-100, 100)), min_size=1)
+    df=dataframes(
+        column("x", dtype=polars.Int64, strategy=integers(-100, 100)),
+        min_size=1,
+        chunked=False if pl_version < Version("0.18.1") else None,
+    )
 )
-def test_transform_function(df: pl.DataFrame, test_funcs):
-    x = pl.col("x")
+def test_transform_function(df: polars.DataFrame, test_funcs):
+    x = polars.col("x")
     transformed_func, original_func = test_funcs
     assert_frame_equal(
         df.select(transformed_func(x).alias("apply")),