1a21aa
From ad34dfa2f983bb3159af8b5780193e0427b505e9 Mon Sep 17 00:00:00 2001
1a21aa
From: Vojtech Trefny <vtrefny@redhat.com>
1a21aa
Date: Thu, 10 Jun 2021 15:01:26 +0200
1a21aa
Subject: [PATCH] Fix/unify importing mock module in tests
1a21aa
1a21aa
mock is available in the unittest module since Python 3.3, we need
1a21aa
to use the old mock module not only with Python 2 but also with
1a21aa
the early versions of Python 3.
1a21aa
---
1a21aa
 tests/action_test.py                         |  5 ++-
1a21aa
 tests/dbus_test.py                           |  9 +++--
1a21aa
 tests/devicefactory_test.py                  |  6 ++--
1a21aa
 tests/devicelibs_test/disk_test.py           |  6 ++--
1a21aa
 tests/devicelibs_test/edd_test.py            |  6 +++-
1a21aa
 tests/devices_test/dependencies_test.py      |  6 ++--
1a21aa
 tests/devices_test/device_methods_test.py    |  8 +++--
1a21aa
 tests/devices_test/device_names_test.py      |  6 ++--
1a21aa
 tests/devices_test/device_properties_test.py |  8 ++---
1a21aa
 tests/devices_test/disk_test.py              |  7 ++--
1a21aa
 tests/devices_test/lvm_test.py               |  7 ++--
1a21aa
 tests/devices_test/partition_test.py         |  6 ++--
1a21aa
 tests/devices_test/tags_test.py              |  6 ++--
1a21aa
 tests/devicetree_test.py                     |  6 ++--
1a21aa
 tests/events_test.py                         |  6 ++--
1a21aa
 tests/formats_test/disklabel_test.py         |  6 ++--
1a21aa
 tests/formats_test/luks_test.py              |  6 ++--
1a21aa
 tests/formats_test/lvmpv_test.py             |  2 --
1a21aa
 tests/formats_test/methods_test.py           |  6 ++--
1a21aa
 tests/formats_test/selinux_test.py           |  6 ++--
1a21aa
 tests/formats_test/swap_test.py              |  2 --
1a21aa
 tests/misc_test.py                           |  6 ++--
1a21aa
 tests/partitioning_test.py                   |  6 ++--
1a21aa
 tests/populator_test.py                      |  6 ++--
1a21aa
 tests/storagetestcase.py                     |  5 ++-
1a21aa
 tests/test_compat.py                         | 38 --------------------
1a21aa
 tests/udev_test.py                           |  6 +++-
1a21aa
 tests/unsupported_disklabel_test.py          |  7 ++--
1a21aa
 tests/util_test.py                           |  6 ++--
1a21aa
 29 files changed, 104 insertions(+), 102 deletions(-)
1a21aa
 delete mode 100644 tests/test_compat.py
1a21aa
1a21aa
diff --git a/tests/action_test.py b/tests/action_test.py
1a21aa
index 38a2e872..f60cf5d7 100644
1a21aa
--- a/tests/action_test.py
1a21aa
+++ b/tests/action_test.py
1a21aa
@@ -1,9 +1,8 @@
1a21aa
-from six import PY3
1a21aa
 import unittest
1a21aa
 
1a21aa
-if PY3:
1a21aa
+try:
1a21aa
     from unittest.mock import Mock
1a21aa
-else:
1a21aa
+except ImportError:
1a21aa
     from mock import Mock
1a21aa
 
1a21aa
 from tests.storagetestcase import StorageTestCase
1a21aa
diff --git a/tests/dbus_test.py b/tests/dbus_test.py
1a21aa
index 293ac073..9bb9102a 100644
1a21aa
--- a/tests/dbus_test.py
1a21aa
+++ b/tests/dbus_test.py
1a21aa
@@ -1,7 +1,10 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
-
1a21aa
 import random
1a21aa
-from six.moves.mock import Mock, patch  # pylint: disable=no-name-in-module,import-error
1a21aa
+
1a21aa
+try:
1a21aa
+    from unittest.mock import patch, Mock
1a21aa
+except ImportError:
1a21aa
+    from mock import patch, Mock
1a21aa
+
1a21aa
 from unittest import TestCase
1a21aa
 
1a21aa
 import dbus
1a21aa
diff --git a/tests/devicefactory_test.py b/tests/devicefactory_test.py
1a21aa
index dc0d6408..93c8bdb7 100644
1a21aa
--- a/tests/devicefactory_test.py
1a21aa
+++ b/tests/devicefactory_test.py
1a21aa
@@ -4,8 +4,10 @@ import unittest
1a21aa
 from decimal import Decimal
1a21aa
 import os
1a21aa
 
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
-from six.moves.mock import patch  # pylint: disable=no-name-in-module,import-error
1a21aa
+try:
1a21aa
+    from unittest.mock import patch
1a21aa
+except ImportError:
1a21aa
+    from mock import patch
1a21aa
 
1a21aa
 import blivet
1a21aa
 
1a21aa
diff --git a/tests/devicelibs_test/disk_test.py b/tests/devicelibs_test/disk_test.py
1a21aa
index e67ef5b1..9cb39951 100644
1a21aa
--- a/tests/devicelibs_test/disk_test.py
1a21aa
+++ b/tests/devicelibs_test/disk_test.py
1a21aa
@@ -1,8 +1,10 @@
1a21aa
 # pylint: skip-file
1a21aa
-import test_compat
1a21aa
+try:
1a21aa
+    from unittest.mock import Mock, patch, sentinel
1a21aa
+except ImportError:
1a21aa
+    from mock import Mock, patch, sentinel
1a21aa
 
1a21aa
 import six
1a21aa
-from six.moves.mock import Mock, patch, sentinel
1a21aa
 import unittest
1a21aa
 
1a21aa
 from blivet.devicelibs import disk as disklib
1a21aa
diff --git a/tests/devicelibs_test/edd_test.py b/tests/devicelibs_test/edd_test.py
1a21aa
index 23d736f4..21bbcffc 100644
1a21aa
--- a/tests/devicelibs_test/edd_test.py
1a21aa
+++ b/tests/devicelibs_test/edd_test.py
1a21aa
@@ -1,5 +1,9 @@
1a21aa
+try:
1a21aa
+    from unittest import mock
1a21aa
+except ImportError:
1a21aa
+    import mock
1a21aa
+
1a21aa
 import unittest
1a21aa
-import mock
1a21aa
 import os
1a21aa
 import inspect
1a21aa
 import logging
1a21aa
diff --git a/tests/devices_test/dependencies_test.py b/tests/devices_test/dependencies_test.py
1a21aa
index c012751d..493d1c9f 100644
1a21aa
--- a/tests/devices_test/dependencies_test.py
1a21aa
+++ b/tests/devices_test/dependencies_test.py
1a21aa
@@ -1,8 +1,6 @@
1a21aa
-# vim:set fileencoding=utf-8
1a21aa
-from six import PY3
1a21aa
-if PY3:
1a21aa
+try:
1a21aa
     from unittest.mock import patch, PropertyMock
1a21aa
-else:
1a21aa
+except ImportError:
1a21aa
     from mock import patch, PropertyMock
1a21aa
 
1a21aa
 import unittest
1a21aa
diff --git a/tests/devices_test/device_methods_test.py b/tests/devices_test/device_methods_test.py
1a21aa
index f00509be..8a70b5bb 100644
1a21aa
--- a/tests/devices_test/device_methods_test.py
1a21aa
+++ b/tests/devices_test/device_methods_test.py
1a21aa
@@ -1,9 +1,11 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
-
1a21aa
 import six
1a21aa
-from six.moves.mock import patch, Mock, PropertyMock  # pylint: disable=no-name-in-module,import-error
1a21aa
 import unittest
1a21aa
 
1a21aa
+try:
1a21aa
+    from unittest.mock import patch, Mock, PropertyMock
1a21aa
+except ImportError:
1a21aa
+    from mock import patch, PropertyMock
1a21aa
+
1a21aa
 from blivet.devices import StorageDevice
1a21aa
 from blivet.devices import DiskDevice, PartitionDevice
1a21aa
 from blivet.devices import LVMVolumeGroupDevice, LVMLogicalVolumeDevice
1a21aa
diff --git a/tests/devices_test/device_names_test.py b/tests/devices_test/device_names_test.py
1a21aa
index 2a66f983..ca44d38c 100644
1a21aa
--- a/tests/devices_test/device_names_test.py
1a21aa
+++ b/tests/devices_test/device_names_test.py
1a21aa
@@ -1,7 +1,9 @@
1a21aa
 # vim:set fileencoding=utf-8
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest.mock import patch
1a21aa
+except ImportError:
1a21aa
+    from mock import patch
1a21aa
 
1a21aa
-from six.moves.mock import patch  # pylint: disable=no-name-in-module,import-error
1a21aa
 import six
1a21aa
 import unittest
1a21aa
 
1a21aa
diff --git a/tests/devices_test/device_properties_test.py b/tests/devices_test/device_properties_test.py
1a21aa
index 240ac088..8928707f 100644
1a21aa
--- a/tests/devices_test/device_properties_test.py
1a21aa
+++ b/tests/devices_test/device_properties_test.py
1a21aa
@@ -1,6 +1,3 @@
1a21aa
-# vim:set fileencoding=utf-8
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
-
1a21aa
 import six
1a21aa
 import unittest
1a21aa
 
1a21aa
@@ -9,7 +6,10 @@ gi.require_version("BlockDev", "2.0")
1a21aa
 
1a21aa
 from gi.repository import BlockDev as blockdev
1a21aa
 
1a21aa
-from six.moves.mock import Mock, patch  # pylint: disable=no-name-in-module,import-error
1a21aa
+try:
1a21aa
+    from unittest.mock import patch, Mock
1a21aa
+except ImportError:
1a21aa
+    from mock import patch, Mock
1a21aa
 
1a21aa
 import blivet
1a21aa
 
1a21aa
diff --git a/tests/devices_test/disk_test.py b/tests/devices_test/disk_test.py
1a21aa
index e9852303..cc8454e1 100644
1a21aa
--- a/tests/devices_test/disk_test.py
1a21aa
+++ b/tests/devices_test/disk_test.py
1a21aa
@@ -1,7 +1,8 @@
1a21aa
 # pylint: skip-file
1a21aa
-import test_compat
1a21aa
-
1a21aa
-from six.moves.mock import patch
1a21aa
+try:
1a21aa
+    from unittest.mock import patch
1a21aa
+except ImportError:
1a21aa
+    from mock import patch
1a21aa
 import unittest
1a21aa
 
1a21aa
 from blivet.devices import DiskDevice
1a21aa
diff --git a/tests/devices_test/lvm_test.py b/tests/devices_test/lvm_test.py
1a21aa
index 670c91c9..f50933c4 100644
1a21aa
--- a/tests/devices_test/lvm_test.py
1a21aa
+++ b/tests/devices_test/lvm_test.py
1a21aa
@@ -1,8 +1,9 @@
1a21aa
-# vim:set fileencoding=utf-8
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest.mock import patch, PropertyMock
1a21aa
+except ImportError:
1a21aa
+    from mock import patch, PropertyMock
1a21aa
 
1a21aa
 import six
1a21aa
-from six.moves.mock import patch, PropertyMock  # pylint: disable=no-name-in-module,import-error
1a21aa
 import unittest
1a21aa
 
1a21aa
 import blivet
1a21aa
diff --git a/tests/devices_test/partition_test.py b/tests/devices_test/partition_test.py
1a21aa
index 0abd88df..4748dafe 100644
1a21aa
--- a/tests/devices_test/partition_test.py
1a21aa
+++ b/tests/devices_test/partition_test.py
1a21aa
@@ -1,5 +1,4 @@
1a21aa
 # vim:set fileencoding=utf-8
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
 
1a21aa
 from collections import namedtuple
1a21aa
 import os
1a21aa
@@ -7,7 +6,10 @@ import six
1a21aa
 import unittest
1a21aa
 import parted
1a21aa
 
1a21aa
-from six.moves.mock import Mock, patch  # pylint: disable=no-name-in-module,import-error
1a21aa
+try:
1a21aa
+    from unittest.mock import patch, Mock
1a21aa
+except ImportError:
1a21aa
+    from mock import patch, Mock
1a21aa
 
1a21aa
 from blivet.devices import DiskFile
1a21aa
 from blivet.devices import PartitionDevice
1a21aa
diff --git a/tests/devices_test/tags_test.py b/tests/devices_test/tags_test.py
1a21aa
index 1edc37f6..49a2d72e 100644
1a21aa
--- a/tests/devices_test/tags_test.py
1a21aa
+++ b/tests/devices_test/tags_test.py
1a21aa
@@ -1,6 +1,8 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest.mock import patch
1a21aa
+except ImportError:
1a21aa
+    from mock import patch
1a21aa
 
1a21aa
-from six.moves.mock import patch  # pylint: disable=no-name-in-module,import-error
1a21aa
 import unittest
1a21aa
 
1a21aa
 from blivet.devices import DiskDevice, FcoeDiskDevice, iScsiDiskDevice, MultipathDevice, StorageDevice, ZFCPDiskDevice
1a21aa
diff --git a/tests/devicetree_test.py b/tests/devicetree_test.py
1a21aa
index 11f8469d..fbf31c77 100644
1a21aa
--- a/tests/devicetree_test.py
1a21aa
+++ b/tests/devicetree_test.py
1a21aa
@@ -1,6 +1,8 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest.mock import patch, Mock, PropertyMock, sentinel
1a21aa
+except ImportError:
1a21aa
+    from mock import patch, Mock, PropertyMock, sentinel
1a21aa
 
1a21aa
-from six.moves.mock import Mock, patch, PropertyMock, sentinel  # pylint: disable=no-name-in-module,import-error
1a21aa
 import os
1a21aa
 import six
1a21aa
 import unittest
1a21aa
diff --git a/tests/events_test.py b/tests/events_test.py
1a21aa
index 5906b4e2..22666f6d 100644
1a21aa
--- a/tests/events_test.py
1a21aa
+++ b/tests/events_test.py
1a21aa
@@ -1,6 +1,8 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest.mock import patch, Mock
1a21aa
+except ImportError:
1a21aa
+    from mock import patch, Mock
1a21aa
 
1a21aa
-from six.moves.mock import Mock, patch  # pylint: disable=no-name-in-module,import-error
1a21aa
 import time
1a21aa
 from unittest import TestCase
1a21aa
 
1a21aa
diff --git a/tests/formats_test/disklabel_test.py b/tests/formats_test/disklabel_test.py
1a21aa
index 0cfa736d..f514a778 100644
1a21aa
--- a/tests/formats_test/disklabel_test.py
1a21aa
+++ b/tests/formats_test/disklabel_test.py
1a21aa
@@ -1,7 +1,9 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest import mock
1a21aa
+except ImportError:
1a21aa
+    import mock
1a21aa
 
1a21aa
 import parted
1a21aa
-from six.moves import mock  # pylint: disable=no-name-in-module,import-error
1a21aa
 import unittest
1a21aa
 
1a21aa
 import blivet
1a21aa
diff --git a/tests/formats_test/luks_test.py b/tests/formats_test/luks_test.py
1a21aa
index be0d50b0..1edbdcb2 100644
1a21aa
--- a/tests/formats_test/luks_test.py
1a21aa
+++ b/tests/formats_test/luks_test.py
1a21aa
@@ -1,6 +1,8 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest.mock import patch
1a21aa
+except ImportError:
1a21aa
+    from mock import patch
1a21aa
 
1a21aa
-from six.moves.mock import patch  # pylint: disable=no-name-in-module,import-error
1a21aa
 import unittest
1a21aa
 
1a21aa
 from blivet.formats.luks import LUKS
1a21aa
diff --git a/tests/formats_test/lvmpv_test.py b/tests/formats_test/lvmpv_test.py
1a21aa
index 792a2f1d..cbd2c419 100644
1a21aa
--- a/tests/formats_test/lvmpv_test.py
1a21aa
+++ b/tests/formats_test/lvmpv_test.py
1a21aa
@@ -1,5 +1,3 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
-
1a21aa
 from blivet.formats.lvmpv import LVMPhysicalVolume
1a21aa
 
1a21aa
 from blivet.size import Size
1a21aa
diff --git a/tests/formats_test/methods_test.py b/tests/formats_test/methods_test.py
1a21aa
index b2674ea7..2743b7db 100644
1a21aa
--- a/tests/formats_test/methods_test.py
1a21aa
+++ b/tests/formats_test/methods_test.py
1a21aa
@@ -1,7 +1,9 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest.mock import patch, sentinel, PropertyMock
1a21aa
+except ImportError:
1a21aa
+    from mock import patch, sentinel, PropertyMock
1a21aa
 
1a21aa
 import six
1a21aa
-from six.moves.mock import patch, sentinel, PropertyMock  # pylint: disable=no-name-in-module,import-error
1a21aa
 import unittest
1a21aa
 
1a21aa
 from blivet.errors import DeviceFormatError
1a21aa
diff --git a/tests/formats_test/selinux_test.py b/tests/formats_test/selinux_test.py
1a21aa
index 02e39011..26df5fe9 100644
1a21aa
--- a/tests/formats_test/selinux_test.py
1a21aa
+++ b/tests/formats_test/selinux_test.py
1a21aa
@@ -1,9 +1,9 @@
1a21aa
 # pylint: disable=unused-import
1a21aa
 import os
1a21aa
-from six import PY3
1a21aa
-if PY3:
1a21aa
+
1a21aa
+try:
1a21aa
     from unittest.mock import patch, ANY
1a21aa
-else:
1a21aa
+except ImportError:
1a21aa
     from mock import patch, ANY
1a21aa
 
1a21aa
 import unittest
1a21aa
diff --git a/tests/formats_test/swap_test.py b/tests/formats_test/swap_test.py
1a21aa
index 56356144..8968ca15 100644
1a21aa
--- a/tests/formats_test/swap_test.py
1a21aa
+++ b/tests/formats_test/swap_test.py
1a21aa
@@ -1,5 +1,3 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
-
1a21aa
 import six
1a21aa
 import unittest
1a21aa
 
1a21aa
diff --git a/tests/misc_test.py b/tests/misc_test.py
1a21aa
index 3c8cf344..10ea1320 100644
1a21aa
--- a/tests/misc_test.py
1a21aa
+++ b/tests/misc_test.py
1a21aa
@@ -1,7 +1,9 @@
1a21aa
 import unittest
1a21aa
 
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
-from six.moves.mock import patch  # pylint: disable=no-name-in-module,import-error
1a21aa
+try:
1a21aa
+    from unittest.mock import patch
1a21aa
+except ImportError:
1a21aa
+    from mock import patch
1a21aa
 
1a21aa
 import blivet
1a21aa
 
1a21aa
diff --git a/tests/partitioning_test.py b/tests/partitioning_test.py
1a21aa
index b7aa5045..9b27f0c0 100644
1a21aa
--- a/tests/partitioning_test.py
1a21aa
+++ b/tests/partitioning_test.py
1a21aa
@@ -1,6 +1,8 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest.mock import patch, Mock
1a21aa
+except ImportError:
1a21aa
+    from mock import patch, Mock
1a21aa
 
1a21aa
-from six.moves.mock import Mock, patch  # pylint: disable=no-name-in-module,import-error
1a21aa
 import six
1a21aa
 import unittest
1a21aa
 
1a21aa
diff --git a/tests/populator_test.py b/tests/populator_test.py
1a21aa
index a7748a9d..2a8532f0 100644
1a21aa
--- a/tests/populator_test.py
1a21aa
+++ b/tests/populator_test.py
1a21aa
@@ -1,7 +1,9 @@
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest.mock import call, patch, sentinel, Mock, PropertyMock
1a21aa
+except ImportError:
1a21aa
+    from mock import call, patch, sentinel, Mock, PropertyMock
1a21aa
 
1a21aa
 import gi
1a21aa
-from six.moves.mock import call, patch, sentinel, Mock, PropertyMock  # pylint: disable=no-name-in-module,import-error
1a21aa
 import six
1a21aa
 import unittest
1a21aa
 
1a21aa
diff --git a/tests/storagetestcase.py b/tests/storagetestcase.py
1a21aa
index 1844dec5..1b856914 100644
1a21aa
--- a/tests/storagetestcase.py
1a21aa
+++ b/tests/storagetestcase.py
1a21aa
@@ -1,6 +1,9 @@
1a21aa
 
1a21aa
 import unittest
1a21aa
-from mock import Mock
1a21aa
+try:
1a21aa
+    from unittest.mock import Mock
1a21aa
+except ImportError:
1a21aa
+    from mock import Mock
1a21aa
 
1a21aa
 import parted
1a21aa
 
1a21aa
diff --git a/tests/test_compat.py b/tests/test_compat.py
1a21aa
deleted file mode 100644
1a21aa
index d0859e24..00000000
1a21aa
--- a/tests/test_compat.py
1a21aa
+++ /dev/null
1a21aa
@@ -1,38 +0,0 @@
1a21aa
-# test_compat.py
1a21aa
-# Python (2 -v- 3) compatibility functions.
1a21aa
-#
1a21aa
-# Copyright (C) 2017  Red Hat, Inc.
1a21aa
-#
1a21aa
-# This copyrighted material is made available to anyone wishing to use,
1a21aa
-# modify, copy, or redistribute it subject to the terms and conditions of
1a21aa
-# the GNU Lesser General Public License v.2, or (at your option) any later
1a21aa
-# version. This program is distributed in the hope that it will be useful,
1a21aa
-# but WITHOUT ANY WARRANTY expressed or implied, including the implied
1a21aa
-# warranties of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See
1a21aa
-# the GNU Lesser General Public License for more details.  You should have
1a21aa
-# received a copy of the GNU Lesser General Public License along with this
1a21aa
-# program; if not, write to the Free Software Foundation, Inc., 51 Franklin
1a21aa
-# Street, Fifth Floor, Boston, MA 02110-1301, USA.  Any Red Hat trademarks
1a21aa
-# that are incorporated in the source code or documentation are not subject
1a21aa
-# to the GNU Lesser General Public License and may only be used or
1a21aa
-# replicated with the express permission of Red Hat, Inc.
1a21aa
-#
1a21aa
-# Red Hat Author(s): David Lehman <dlehman@redhat.com>
1a21aa
-#
1a21aa
-
1a21aa
-import six as _six
1a21aa
-
1a21aa
-mock_move = _six.MovedModule('mock', 'mock', 'unittest.mock')
1a21aa
-
1a21aa
-
1a21aa
-def add_move(mod):
1a21aa
-    _six.add_move(mod)
1a21aa
-    # https://bitbucket.org/gutworth/six/issues/116/enable-importing-from-within-custom
1a21aa
-    _six._importer._add_module(mod, "moves." + mod.name)
1a21aa
-
1a21aa
-
1a21aa
-def setup():
1a21aa
-    add_move(mock_move)
1a21aa
-
1a21aa
-
1a21aa
-setup()
1a21aa
diff --git a/tests/udev_test.py b/tests/udev_test.py
1a21aa
index f9b10620..569a144e 100644
1a21aa
--- a/tests/udev_test.py
1a21aa
+++ b/tests/udev_test.py
1a21aa
@@ -1,6 +1,10 @@
1a21aa
 
1a21aa
 import unittest
1a21aa
-import mock
1a21aa
+
1a21aa
+try:
1a21aa
+    from unittest import mock
1a21aa
+except ImportError:
1a21aa
+    import mock
1a21aa
 
1a21aa
 from udev_data import raid_data
1a21aa
 
1a21aa
diff --git a/tests/unsupported_disklabel_test.py b/tests/unsupported_disklabel_test.py
1a21aa
index f5b24779..38055333 100644
1a21aa
--- a/tests/unsupported_disklabel_test.py
1a21aa
+++ b/tests/unsupported_disklabel_test.py
1a21aa
@@ -1,7 +1,8 @@
1a21aa
-# vim:set fileencoding=utf-8
1a21aa
-import test_compat  # pylint: disable=unused-import
1a21aa
+try:
1a21aa
+    from unittest.mock import patch, sentinel, DEFAULT
1a21aa
+except ImportError:
1a21aa
+    from mock import patch, sentinel, DEFAULT
1a21aa
 
1a21aa
-from six.moves.mock import patch, sentinel, DEFAULT  # pylint: disable=no-name-in-module,import-error
1a21aa
 import six
1a21aa
 import unittest
1a21aa
 
1a21aa
diff --git a/tests/util_test.py b/tests/util_test.py
1a21aa
index 853b6166..b4f82c1b 100644
1a21aa
--- a/tests/util_test.py
1a21aa
+++ b/tests/util_test.py
1a21aa
@@ -1,7 +1,9 @@
1a21aa
 # pylint: skip-file
1a21aa
-import test_compat
1a21aa
+try:
1a21aa
+    from unittest import mock
1a21aa
+except ImportError:
1a21aa
+    import mock
1a21aa
 
1a21aa
-from six.moves import mock
1a21aa
 import os
1a21aa
 import six
1a21aa
 import tempfile
1a21aa
-- 
1a21aa
2.31.1
1a21aa