diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py
index 8b45c2148504..a2b02b56692c 100644
--- a/tests/models/test_attention_processor.py
+++ b/tests/models/test_attention_processor.py
@@ -1,6 +1,5 @@
 import importlib.metadata
 import tempfile
-import unittest
 
 import numpy as np
 import pytest
@@ -13,7 +12,7 @@
 from ..testing_utils import torch_device
 
 
-class AttnAddedKVProcessorTests(unittest.TestCase):
+class TestAttnAddedKVProcessor:
     def get_constructor_arguments(self, only_cross_attention: bool = False):
         query_dim = 10
 
@@ -55,8 +54,8 @@ def test_only_cross_attention(self):
         constructor_args = self.get_constructor_arguments(only_cross_attention=False)
         attn = Attention(**constructor_args)
 
-        self.assertTrue(attn.to_k is not None)
-        self.assertTrue(attn.to_v is not None)
+        assert attn.to_k is not None
+        assert attn.to_v is not None
 
         forward_args = self.get_forward_arguments(
             query_dim=constructor_args["query_dim"], added_kv_proj_dim=constructor_args["added_kv_proj_dim"]
@@ -71,8 +70,8 @@ def test_only_cross_attention(self):
         constructor_args = self.get_constructor_arguments(only_cross_attention=True)
         attn = Attention(**constructor_args)
 
-        self.assertTrue(attn.to_k is None)
-        self.assertTrue(attn.to_v is None)
+        assert attn.to_k is None
+        assert attn.to_v is None
 
         forward_args = self.get_forward_arguments(
             query_dim=constructor_args["query_dim"], added_kv_proj_dim=constructor_args["added_kv_proj_dim"]
@@ -80,10 +79,10 @@ def test_only_cross_attention(self):
 
         only_cross_attn_out = attn(**forward_args)
 
-        self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all())
+        assert (only_cross_attn_out != self_and_cross_attn_out).all()
 
 
-class DeprecatedAttentionBlockTests(unittest.TestCase):
+class TestDeprecatedAttentionBlock:
     @pytest.fixture(scope="session")
     def is_dist_enabled(pytestconfig):
         return pytestconfig.getoption("dist") == "loadfile"
@@ -131,5 +130,5 @@ def test_conversion_when_using_device_map(self):
             output_type="np",
         ).images
 
-        self.assertTrue(np.allclose(pre_conversion, conversion, atol=1e-3))
-        self.assertTrue(np.allclose(conversion, after_conversion, atol=1e-3))
+        assert np.allclose(pre_conversion, conversion, atol=1e-3)
+        assert np.allclose(conversion, after_conversion, atol=1e-3)