From 9b761f04cdfab8108ab3896848132fec9d189f8d Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 23 Jul 2024 16:47:56 +0000 Subject: [PATCH] Fix test --- python/cudf_polars/tests/test_config.py | 50 +++++++++++++------------ 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 70beee9873e..5b4bba55552 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -3,30 +3,32 @@ from __future__ import annotations -import os +import pytest import polars as pl -from cudf_polars.dsl.expr import Expr -import cudf._lib.pylibcudf as plc - - -class NotImplementedExpr(Expr): - def do_evaluate(self, df, *, context, mapping): - raise NotImplementedError("Example Unsupported cuDF operation") - - -def invoke_unimplemented_op(df): - # TODO: does not work - # df.select(NotImplementedExpr(plc.DataType(plc.TypeId.INT64))) - pass - - -def test_polars_verbose_warns(): - old_val = os.environ.get("POLARS_VERBOSE") - os.environ["POLARS_VERBOSE"] = "1" - df = pl.DataFrame({"a": [1, 2, 3]}).lazy() - - invoke_unimplemented_op(df) - - os.environ["POLARS_VERBOSE"] = old_val +from cudf_polars.dsl.ir import IR +from cudf_polars.testing.asserts import ( + assert_gpu_result_equal, + assert_ir_translation_raises, +) + + +def test_polars_verbose_warns(monkeypatch): + def raise_unimplemented(self): + raise NotImplementedError("We don't support this") + + monkeypatch.setattr(IR, "__post_init__", raise_unimplemented) + q = pl.LazyFrame({}) + # Ensure that things raise + assert_ir_translation_raises(q, NotImplementedError) + with ( + pl.Config(verbose=True), + pytest.raises(pl.exceptions.ComputeError), + pytest.warns( + pl.exceptions.PerformanceWarning, + match="Query execution with GPU not supported", + ), + ): + # And ensure that collecting issues the correct warning. + assert_gpu_result_equal(q)