RL-2204: Update test and type

(Issue RL-2204 - Improve support for categorical plot)

Author: EnliteAI Bot

maze/core/log_events/log_create_figure_functions.py

@@ -27,7 +27,7 @@ def create_binary_plot(value: Union[List[Tuple[np.ndarray, int]], List[int], Lis
 
 
 def create_categorical_plot(
-    value: Union[List[Tuple[int, int]], List[int], List[float], List[np.integer], List[np.floating]], **kwargs
+    value: Union[List[Tuple[int, int]], List[int], List[float], np.ndarray], **kwargs
 ) -> plt.Figure:
     """Checks the type of value and calls the correct plotting function accordingly.
 

maze/test/core/log_events/test_log_create_figure_functions.py

@@ -15,18 +15,23 @@ def test_create_histogram():
 
 
 def test_create_categorical_plot():
-    """ unit test """
+    """ Test creation of categorical plot with integer, tuples and float. """
+
+    # test with integer
     values = [int(i) for i in np.random.randint(low=0, high=10, size=100)]
     fig = create_categorical_plot(values)
     plt.close(fig)
 
+    # test with tuples
     values = [(v, 10) for v in np.random.randint(low=0, high=10, size=100)]
     fig = create_categorical_plot(values)
     plt.close(fig)
 
-    with raises(NotImplementedError):
-        values = np.random.random(100)
-        create_categorical_plot(values)
+    # test with np.float
+    values = np.random.random(100)
+    create_categorical_plot(values)
+    plt.close(fig)
+
 
 
 def test_create_multi_binary_plot():