From 21dc13b8def4d7395c2c83c0b63bf310354117e2 Mon Sep 17 00:00:00 2001
From: Raoul Schram <r.d.schram@uu.nl>
Date: Wed, 6 Mar 2024 13:50:48 +0100
Subject: [PATCH] Fix issue with polars update

---
 metasyncontrib/disclosure/utils.py |  6 +-----
 tests/test_other_dist.py           | 13 ++++++++-----
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/metasyncontrib/disclosure/utils.py b/metasyncontrib/disclosure/utils.py
index 9416deb..68ed5d0 100644
--- a/metasyncontrib/disclosure/utils.py
+++ b/metasyncontrib/disclosure/utils.py
@@ -61,11 +61,7 @@ def _create_subsample( # pylint: disable=too-many-locals
         sub_values = []
         for block in block_values:
             mean_time = pl.Series(block).dt.cast_time_unit("us").mean()
-            assert mean_time is not None
-            sec_since_1970 = mean_time / 1e6
-            sub_values.append(
-                dt.datetime.utcfromtimestamp(0) + dt.timedelta(seconds=sec_since_1970)
-            )
+            sub_values.append(mean_time)
     return sub_values, dominance
 
 
diff --git a/tests/test_other_dist.py b/tests/test_other_dist.py
index 2f60f71..2efd121 100644
--- a/tests/test_other_dist.py
+++ b/tests/test_other_dist.py
@@ -1,3 +1,5 @@
+import datetime as dt
+
 import numpy as np
 import polars as pl
 from metasyn.distribution.categorical import MultinoulliDistribution
@@ -14,20 +16,21 @@
 
 
 @mark.parametrize(
-    "class_norm,class_disc",
+    "class_norm,class_disc,input_type",
     [
-        (DateUniformDistribution, DisclosureDate),
-        (DateTimeUniformDistribution, DisclosureDateTime),
-        (TimeUniformDistribution, DisclosureTime),
+        (DateUniformDistribution, DisclosureDate, dt.date),
+        (DateTimeUniformDistribution, DisclosureDateTime, dt.datetime),
+        (TimeUniformDistribution, DisclosureTime, dt.time),
     ],
 )
-def test_datetime(class_norm, class_disc):
+def test_datetime(class_norm, class_disc, input_type):
     dist_norm = class_norm.default_distribution()
     series = pl.Series([dist_norm.draw() for _ in range(100)])
     dist_norm = class_norm.fit(series)
     dist_disc = class_disc.fit(series)
     assert dist_norm.lower < dist_disc.lower
     assert dist_norm.upper > dist_disc.upper
+    assert isinstance(dist_norm.draw(), input_type)
     if not isinstance(dist_norm, DateUniformDistribution):
         assert dist_norm.precision == dist_disc.precision