tests

cvxgrp · Dec 22, 2023 · 8d877ae · 8d877ae
1 parent 3a35888
commit 8d877ae
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/tests/test_dataclass.py b/tests/test_dataclass.py
@@ -2,22 +2,23 @@
 
 import numpy as np
 import pandas as pd
+import polars as pl
 
 from cvx.bson.dataclass import Data
 
 
 @dataclass(frozen=True)
 class Maffay(Data):
     x: pd.DataFrame
-    y: pd.DataFrame
+    y: pl.DataFrame
     z: np.array
 
 
 def test_conversion(tmp_path):
     matrix = np.random.rand(5, 2)
 
     x = pd.DataFrame(data=matrix)
-    y = pd.DataFrame(data=2 * matrix)
+    y = pl.DataFrame(data=2 * matrix)
     z = matrix
 
     data = Maffay(x=x, y=y, z=z)

diff --git a/tests/test_polars.py → tests/test_encode_decode.py b/tests/test_polars.py → tests/test_encode_decode.py
diff --git a/tests/test_pandas.py b/tests/test_pandas.py
@@ -28,6 +28,9 @@ def assert_equal(obj1, obj2):
     if isinstance(obj1, np.ndarray):
         np.testing.assert_array_equal(obj1, obj2)
 
+    if isinstance(obj1, pl.DataFrame):
+        assert obj1.equals(obj2)
+
 
 def test_roundtrip(data):
     """