updated test_dataset_generator asserts and get_col_ordered_column test

capitalone · Jul 26, 2023 · 4cfc611 · 4cfc611
1 parent 189a969
commit 4cfc611
Showing 1 changed file with 51 additions and 7 deletions.
diff --git a/tests/test_dataset_generator.py b/tests/test_dataset_generator.py
@@ -63,7 +63,7 @@ def test_get_ordered_column_datetime_ascending(self):
         ordered_data = ordered_data[:, 0]
         output_data = dg.get_ordered_column(data, "datetime", "ascending")
 
-        self.assertTrue(np.array_equal(output_data, ordered_data))
+        np.testing.assert_array_equal(output_data, ordered_data)
 
     def test_get_ordered_column_datetime_descending(self):
         data = dategen.random_datetimes(
@@ -108,7 +108,7 @@ def test_get_ordered_column_datetime_descending(self):
         ordered_data = ordered_data[:, 0]
         output_data = dg.get_ordered_column(data, "datetime", "descending")
 
-        self.assertTrue(np.array_equal(output_data, ordered_data))
+        np.testing.assert_array_equal(output_data, ordered_data)
 
     def test_get_ordered_column_custom_datetime_ascending(self):
         custom_date_format = ["%Y %m %d"]
@@ -148,7 +148,7 @@ def test_get_ordered_column_custom_datetime_ascending(self):
         ordered_data = ordered_data[:, 0]
         output_data = dg.get_ordered_column(data, "datetime", "ascending")
 
-        self.assertTrue(np.array_equal(output_data, ordered_data))
+        np.testing.assert_array_equal(output_data, ordered_data)
 
     def test_get_ordered_column_custom_datetime_descending(self):
         custom_date_format = ["%Y %m %d"]
@@ -188,7 +188,7 @@ def test_get_ordered_column_custom_datetime_descending(self):
         ordered_data = ordered_data[:, 0]
         output_data = dg.get_ordered_column(data, "datetime", "descending")
 
-        self.assertTrue(np.array_equal(output_data, ordered_data))
+        np.testing.assert_array_equal(output_data, ordered_data)
 
     def test_get_ordered_column(self):
 
@@ -198,19 +198,63 @@ def test_get_ordered_column(self):
                 "float": np.array([5.0, 4.0, 3.0, 2.0, 1.0]),
                 "string": np.array(["abcde", "bcdea", "cdeab", "deabc", "eabcd"]),
                 "categorical": np.array(["E", "D", "C", "B", "A"]),
+                "datetime": np.array(
+                    [
+                        [
+                            "September 27 2018 18:24:03",
+                            datetime.strptime(
+                                "September 27 2018 18:24:03", self.date_format_list[0]
+                            ),
+                        ],
+                        [
+                            "March 11 2016 15:15:39",
+                            datetime.strptime(
+                                "March 11 2016 15:15:39", self.date_format_list[0]
+                            ),
+                        ],
+                        [
+                            "March 13 2010 17:18:44",
+                            datetime.strptime(
+                                "March 13 2010 17:18:44", self.date_format_list[0]
+                            ),
+                        ],
+                        [
+                            "August 19 2008 16:53:49",
+                            datetime.strptime(
+                                "August 19 2008 16:53:49", self.date_format_list[0]
+                            ),
+                        ],
+                        [
+                            "October 02 2006 22:34:32",
+                            datetime.strptime(
+                                "October 02 2006 22:34:32", self.date_format_list[0]
+                            ),
+                        ],
+                    ]
+                ),
             }
         )
+
         ordered_data = [
             np.array([1, 2, 3, 4, 5]),
             np.array([1.0, 2.0, 3.0, 4.0, 5.0]),
             np.array(["abcde", "bcdea", "cdeab", "deabc", "eabcd"]),
             np.array(["A", "B", "C", "D", "E"]),
+            np.array(
+                [
+                    "October 02 2006 22:34:32",
+                    "August 19 2008 16:53:49",
+                    "March 13 2010 17:18:44",
+                    "March 11 2016 15:15:39",
+                    "September 27 2018 18:24:03",
+                ]
+            ),
         ]
-        ordered_data = np.asarray(ordered_data)
+        ordered_data = np.array(ordered_data, dtype=object)
 
         output_data = []
         for data_type in data.keys():
             output_data.append(dg.get_ordered_column(data[data_type], data_type))
-        output_data = np.asarray(output_data)
+        output_data = np.array(output_data)
 
-        self.assertTrue(np.array_equal(output_data, ordered_data))
+        np.testing.assert_array_equal(output_data, ordered_data)