From a68d5959780755475173badd97a06f79c551fd86 Mon Sep 17 00:00:00 2001
From: Alexandre Detiste <alexandre.detiste@gmail.com>
Date: Wed, 17 Apr 2024 10:18:15 +0200
Subject: [PATCH] prefer newer unittest.mock when available

---
 tests/test_progressbar.py           | 6 +++++-
 tests/test_query_device_capacity.py | 6 +++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/tests/test_progressbar.py b/tests/test_progressbar.py
index 304e14e..13b2427 100644
--- a/tests/test_progressbar.py
+++ b/tests/test_progressbar.py
@@ -28,10 +28,14 @@
 Test the progressbar 'FileTransferSpeed' integration
 """
 
+try:
+    import unittest.mock
+except ImportError:
+    import mock
+
 from . import TestCase
 import bitmath
 from bitmath.integrations.bmprogressbar import BitmathFileTransferSpeed
-import mock
 import progressbar
 
 
diff --git a/tests/test_query_device_capacity.py b/tests/test_query_device_capacity.py
index 64ca99e..4ddda41 100644
--- a/tests/test_query_device_capacity.py
+++ b/tests/test_query_device_capacity.py
@@ -30,9 +30,13 @@
 
 from . import TestCase
 import bitmath
-import mock
 import struct
 
+try:
+    import unittest.mock
+except ImportError:
+    import mock
+
 try:
     # Python 3.3+
     from contextlib import ExitStack, contextmanager