From 8d877ae9bf578f1815e63f1561d6a2af7666852b Mon Sep 17 00:00:00 2001
From: Thomas Schmelzer <thomas.schmelzer@gmail.com>
Date: Fri, 22 Dec 2023 17:31:57 +0400
Subject: [PATCH] tests

---
 tests/test_dataclass.py                         | 5 +++--
 tests/{test_polars.py => test_encode_decode.py} | 0
 tests/test_pandas.py                            | 3 +++
 3 files changed, 6 insertions(+), 2 deletions(-)
 rename tests/{test_polars.py => test_encode_decode.py} (100%)

diff --git a/tests/test_dataclass.py b/tests/test_dataclass.py
index 87a4329c..f37cad2a 100644
--- a/tests/test_dataclass.py
+++ b/tests/test_dataclass.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 import pandas as pd
+import polars as pl
 
 from cvx.bson.dataclass import Data
 
@@ -9,7 +10,7 @@
 @dataclass(frozen=True)
 class Maffay(Data):
     x: pd.DataFrame
-    y: pd.DataFrame
+    y: pl.DataFrame
     z: np.array
 
 
@@ -17,7 +18,7 @@ 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 b/tests/test_encode_decode.py
similarity index 100%
rename from tests/test_polars.py
rename to tests/test_encode_decode.py
diff --git a/tests/test_pandas.py b/tests/test_pandas.py
index 8e848b33..b8bd5096 100644
--- 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):
     """