From de73d39da47dcaab7afef76d1758600d46dfa37c Mon Sep 17 00:00:00 2001
From: Niru Maheswaranathan <niru@hey.com>
Date: Tue, 20 May 2025 04:41:43 -0700
Subject: [PATCH] Add unit tests for chart utilities and cmat

---
 tests/test_images.py | 11 +++++++++
 tests/test_utils.py  | 55 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/tests/test_images.py b/tests/test_images.py
index cabef65..8c6c05d 100644
--- a/tests/test_images.py
+++ b/tests/test_images.py
@@ -16,3 +16,14 @@ def test_img_corr_mode():
     # Colorbar should have been added
     assert len(fig.axes) == 2
     plt.close(fig)
+
+def test_cmat_labels_and_colorbar():
+    data = np.array([[0.0, 1.0], [1.0, 0.0]])
+    fig, ax = plt.subplots()
+    cb, returned_ax = images.cmat(data, labels=["a", "b"], cbar=True, fig=fig, ax=ax)
+
+    assert returned_ax is ax
+    assert [tick.get_text() for tick in ax.get_xticklabels()] == ["a", "b"]
+    assert [tick.get_text() for tick in ax.get_yticklabels()] == ["a", "b"]
+    assert len(fig.axes) == 2
+    plt.close(fig)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index ee9c161..5d3be0e 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -51,3 +51,58 @@ def test_noticks():
     assert len(ax.get_yticks()) == 0
 
     plt.close(fig)
+
+def test_get_bounds_spines():
+    fig, ax = plt.subplots()
+    ax.plot([0, 1], [0, 1])
+    ax.spines["bottom"].set_bounds(0, 1)
+    ax.spines["left"].set_bounds(-1, 2)
+
+    assert cu.get_bounds("x", ax=ax) == (0, 1)
+    assert cu.get_bounds("y", ax=ax) == (-1, 2)
+    plt.close(fig)
+
+
+def test_get_bounds_label_fallback():
+    fig, ax = plt.subplots()
+    ax.plot([0, 1, 2, 3], [0, 1, 2, 3])
+    ax.set_xticks([0, 1, 2, 3])
+    ax.set_xticklabels(["", "a", "", "b"])
+
+    assert cu.get_bounds("x", ax=ax) == (1, 3)
+    plt.close(fig)
+
+
+def test_breathe_adds_padding_and_hides_spines():
+    fig, ax = plt.subplots()
+    ax.plot([0, 1], [0, 1])
+    old_xlim = ax.get_xlim()
+    old_ylim = ax.get_ylim()
+
+    cu.breathe(ax=ax)
+
+    new_xlim = ax.get_xlim()
+    new_ylim = ax.get_ylim()
+
+    assert new_xlim[0] < old_xlim[0] and new_xlim[1] > old_xlim[1]
+    assert new_ylim[0] < old_ylim[0] and new_ylim[1] > old_ylim[1]
+
+    assert not ax.spines["top"].get_visible()
+    assert not ax.spines["right"].get_visible()
+    assert ax.spines["bottom"].get_bounds() == old_xlim
+    assert ax.spines["left"].get_bounds() == old_ylim
+    plt.close(fig)
+
+
+def test_xclamp_yclamp():
+    fig, ax = plt.subplots()
+    ax.plot([0, 1, 2, 3, 4], [0, 1, 2, 3, 4])
+
+    cu.xclamp(x0=0.5, x1=3.4, dt=1.0, ax=ax)
+    assert ax.get_xlim() == (0.0, 4.0)
+    assert list(ax.get_xticks()) == [0.0, 1.0, 2.0, 3.0, 4.0]
+
+    cu.yclamp(y0=0.2, y1=3.5, dt=1.0, ax=ax)
+    assert ax.get_ylim() == (0.0, 4.0)
+    assert list(ax.get_yticks()) == [0.0, 1.0, 2.0, 3.0, 4.0]
+    plt.close(fig)