diff --git a/tests/test_channel_metrics/test_completely_bounded_trace_norm.py b/tests/test_channel_metrics/test_completely_bounded_trace_norm.py
index 86c920be0..becd6b373 100644
--- a/tests/test_channel_metrics/test_completely_bounded_trace_norm.py
+++ b/tests/test_channel_metrics/test_completely_bounded_trace_norm.py
@@ -1,5 +1,6 @@
 """Tests for completely_bounded_trace_norm."""
 import numpy as np
+import pytest
 
 from toqito.channel_metrics import completely_bounded_trace_norm
 from toqito.channel_ops import kraus_to_choi
@@ -9,7 +10,13 @@
 def test_cb_trace_norm_quantum_channel():
     """The diamond norm of a quantum channel is 1."""
     phi = dephasing(2)
-    np.testing.assert_equal(completely_bounded_trace_norm(phi), 1)
+    assert completely_bounded_trace_norm(phi) == 1
+
+
+def test_cb_trace_norm_CP():
+    """Test for the diamond norm of a CP map."""
+    non_normalized_depolarizing_array = np.eye(4)
+    assert completely_bounded_trace_norm(non_normalized_depolarizing_array) == 4.0
 
 
 def test_cb_trace_norm_unitaries_channel():
@@ -19,17 +26,14 @@ def test_cb_trace_norm_unitaries_channel():
     lam, _ = np.linalg.eig(U)
     dist = np.abs(lam[:, None] - lam[None, :])  # all to all distance
     diameter = np.max(dist)
-    np.testing.assert_equal(
-        np.isclose(completely_bounded_trace_norm(phi), diameter, atol=1e-3), True
-    )
+    assert np.isclose(completely_bounded_trace_norm(phi), diameter, atol=1e-3)
 
 
 def test_cb_trace_norm_non_square():
     """Non-square inputs for cb trace norm."""
-    with np.testing.assert_raises(ValueError):
+    with pytest.raises(
+        ValueError,
+        match="The input and output spaces of the superoperator phi must both be square.",
+    ):
         phi = np.array([[1, 2, 3], [4, 5, 6]])
         completely_bounded_trace_norm(phi)
-
-
-if __name__ == "__main__":
-    np.testing.run_module_suite()