diff --git a/leafmap/colormaps.py b/leafmap/colormaps.py
index 6eea655998..cb50f2a874 100644
--- a/leafmap/colormaps.py
+++ b/leafmap/colormaps.py
@@ -46,7 +46,7 @@ def get_palette(
     cmap_name: Optional[str] = None,
     n_class: Optional[int] = None,
     hashtag: Optional[bool] = False,
-) -> List:
+) -> List[str]:
     """Get a palette from a matplotlib colormap. See the list of colormaps at https://matplotlib.org/stable/tutorials/colors/colormaps.html.
 
     Args:
@@ -81,7 +81,7 @@ def get_colorbar(
     height: Optional[float] = 0.4,
     orientation: Optional[str] = "horizontal",
     discrete: Optional[bool] = False,
-):
+) -> None:
     """Creates a colorbar based on custom colors.
 
     Args:
@@ -135,7 +135,7 @@ def create_colormap(
     show_name: Optional[bool] = False,
     font_size: Optional[int] = 12,
     **kwargs
-):
+) -> None:
     """Plot a matplotlib colormap.
 
     Args:
@@ -204,7 +204,7 @@ def plot_colormap(
     show_name: Optional[bool] = False,
     font_size: Optional[int] = 12,
     **kwargs
-):
+) -> None:
     """Plot a matplotlib colormap.
 
     Args:
@@ -264,7 +264,7 @@ def plot_colormaps(
     height: Optional[float] = 0.4,
     return_fig: Optional[bool] = False,
     **kwargs
-):
+) -> None:
     """Plot all available colormaps.
 
     Args:
diff --git a/tests/test_colormaps.py b/tests/test_colormaps.py
index e5253de539..544bc098c6 100644
--- a/tests/test_colormaps.py
+++ b/tests/test_colormaps.py
@@ -15,6 +15,11 @@ class TestColormaps(unittest.TestCase):
 
     def setUp(self):
         """Set up test fixtures, if any."""
+        # self.__palette_dict: dict = {
+        #     "dem": ["#000000", "#FFFFFF"],
+        #     "ndvi": ["#00FF00", "#FF0000"],
+        #     "ndwi": ["#0000FF", "#00FFFF"],
+        # }
 
     def tearDown(self):
         """Tear down test fixtures, if any."""
@@ -32,6 +37,16 @@ def test_get_palette(self):
         # with self.assertRaises(ValueError):
         #     cm.get_palette("not_a_palette")
 
+    def test_invalid_colormap(self):
+        with self.assertRaises(KeyError):
+            cm.get_palette("invd", 5, True)
+
+    def test_colormap_with_n_class(self):
+        res = cm.get_palette("viridis", 5, False)
+        self.assertEqual(len(res), 5)
+        for color in res:
+            self.assertFalse(color.startswith("#"))
+
     # @patch("matplotlib.pyplot.show")
     # def test_get_colorbar(self, mock_show):
     #     """Test the get_colorbar function."""