Update test cases for GT.tab_style()

tests/test_tab_create_modify.py

@@ -13,6 +13,11 @@ def gt():
     return GT(pd.DataFrame({"x": [1, 2], "y": [4, 5]}))
 
 
+@pytest.fixture
+def gt2():
+    return GT(pl.DataFrame({"x": [1, 2], "y": [4, 5]}))
+
+
 def test_tab_style(gt: GT):
     style = CellStyleFill(color="blue")
     new_gt = tab_style(gt, style, LocBody(["x"], [0]))
@@ -34,30 +39,6 @@ def test_tab_style_multiple_columns(gt: GT):
     assert new_gt._styles[0].styles[0] is style
 
 
-def test_tab_style_loc_body_mask():
-    gt = GT(pl.DataFrame({"x": [1, 2], "y": [4, 5]}))
-    style = CellStyleFill(color="blue")
-    new_gt = tab_style(gt, style, LocBody(mask=cs.numeric().gt(1.5)))
-
-    assert len(gt._styles) == 0
-    assert len(new_gt._styles) == 3
-
-    xy_0y, xy_1x, xy_1y = new_gt._styles
-
-    assert xy_0y.styles[0] is style
-    assert xy_1x.styles[0] is style
-    assert xy_1y.styles[0] is style
-
-    assert xy_0y.rownum == 0
-    assert xy_0y.colname == "y"
-
-    assert xy_1x.rownum == 1
-    assert xy_1x.colname == "x"
-
-    assert xy_1y.rownum == 1
-    assert xy_1y.colname == "y"
-
-
 def test_tab_style_google_font(gt: GT):
     new_gt = tab_style(
         gt,
@@ -96,3 +77,51 @@ def test_tab_style_font_from_column():
 
     assert rendered_html.find('<td style="font-family: Helvetica;" class="gt_row gt_right">1</td>')
     assert rendered_html.find('<td style="font-family: Courier;" class="gt_row gt_right">2</td>')
+
+
+def test_tab_style_loc_body_mask(gt2: GT):
+    style = CellStyleFill(color="blue")
+    new_gt = tab_style(gt2, style, LocBody(mask=cs.numeric().gt(1.5)))
+
+    assert len(gt2._styles) == 0
+    assert len(new_gt._styles) == 3
+
+    xy_0y, xy_1x, xy_1y = new_gt._styles
+
+    assert xy_0y.styles[0] is style
+    assert xy_1x.styles[0] is style
+    assert xy_1y.styles[0] is style
+
+    assert xy_0y.rownum == 0
+    assert xy_0y.colname == "y"
+
+    assert xy_1x.rownum == 1
+    assert xy_1x.colname == "x"
+
+    assert xy_1y.rownum == 1
+    assert xy_1y.colname == "y"
+
+
+def test_tab_style_loc_body_raises(gt2: GT):
+    style = CellStyleFill(color="blue")
+    mask = cs.numeric().gt(1.5)
+    err_msg = "Cannot specify the `mask` argument along with `columns` or `rows` in `loc.body()`."
+
+    with pytest.raises(ValueError) as exc_info:
+        tab_style(gt2, style, LocBody(columns=["x"], mask=mask))
+    assert err_msg in exc_info.value.args[0]
+
+    with pytest.raises(ValueError) as exc_info:
+        tab_style(gt2, style, LocBody(rows=[0], mask=mask))
+
+    assert err_msg in exc_info.value.args[0]
+
+
+def test_tab_style_loc_body_mask_not_polars_expression_raises(gt2: GT):
+    style = CellStyleFill(color="blue")
+    mask = "fake expression"
+    err_msg = "Only Polars expressions can be passed to the `mask` argument."
+
+    with pytest.raises(ValueError) as exc_info:
+        tab_style(gt2, style, LocBody(mask=mask))
+    assert err_msg in exc_info.value.args[0]