refactor ionice() on Linux

2 files changed, 36 insertions, 59 deletions

diff --git a/psutil/_pslinux.py b/psutil/_pslinux.py
index ecba4139..7c9c669f 100644
--- a/psutil/_pslinux.py
+++ b/psutil/_pslinux.py
@@ -41,7 +41,6 @@ from ._common import supports_ipv6
 from ._common import usage_percent
 from ._compat import b
 from ._compat import basestring
-from ._compat import long
 from ._compat import PY3
 
 if sys.version_info >= (3, 4):
@@ -1987,36 +1986,16 @@ class Process(object):
 
         @wrap_exceptions
         def ionice_set(self, ioclass, value):
-            if value is not None:
-                if not PY3 and not isinstance(value, (int, long)):
-                    msg = "value argument is not an integer (gor %r)" % value
-                    raise TypeError(msg)
-                if not 0 <= value <= 7:
-                    raise ValueError(
-                        "value argument range expected is between 0 and 7")
-
-            if ioclass in (IOPRIO_CLASS_NONE, None):
-                if value:
-                    msg = "can't specify value with IOPRIO_CLASS_NONE " \
-                          "(got %r)" % value
-                    raise ValueError(msg)
-                ioclass = IOPRIO_CLASS_NONE
+            if value is None:
                 value = 0
-            elif ioclass == IOPRIO_CLASS_IDLE:
-                if value:
-                    msg = "can't specify value with IOPRIO_CLASS_IDLE " \
-                          "(got %r)" % value
-                    raise ValueError(msg)
-                value = 0
-            elif ioclass in (IOPRIO_CLASS_RT, IOPRIO_CLASS_BE):
-                if value is None:
-                    # TODO: add comment explaining why this is 4 (?)
-                    value = 4
-            else:
-                # otherwise we would get OSError(EVINAL)
-                raise ValueError("invalid ioclass argument %r" % ioclass)
-
-            return cext.proc_ioprio_set(self.pid, ioclass, value)
+            if value and ioclass == IOPRIO_CLASS_IDLE:
+                raise ValueError("IOPRIO_CLASS_IDLE accepts no value")
+            try:
+                return cext.proc_ioprio_set(self.pid, ioclass, value)
+            except OSError as err:
+                if err.errno == errno.EINVAL and value > 7:
+                    raise ValueError("value not in 0-7 range")
+                raise
 
     if HAS_PRLIMIT:
 
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index ae231e35..d264ce6d 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -352,38 +352,36 @@ class TestProcess(unittest.TestCase):
     @unittest.skipIf(not LINUX, "linux only")
     def test_ionice_linux(self):
         p = psutil.Process()
-        self.assertEqual(tuple(p.ionice()), (psutil.IOPRIO_CLASS_NONE, 0))
-
+        self.assertEqual(p.ionice()[0], psutil.IOPRIO_CLASS_NONE)
         self.assertEqual(psutil.IOPRIO_CLASS_NONE, 0)
-        self.assertEqual(psutil.IOPRIO_CLASS_RT, 1)
-        self.assertEqual(psutil.IOPRIO_CLASS_BE, 2)
-        self.assertEqual(psutil.IOPRIO_CLASS_IDLE, 3)
-        if enum is not None:
-            self.assertIsInstance(p.ionice()[0], enum.IntEnum)
-
+        self.assertEqual(psutil.IOPRIO_CLASS_RT, 1)  # high
+        self.assertEqual(psutil.IOPRIO_CLASS_BE, 2)  # normal
+        self.assertEqual(psutil.IOPRIO_CLASS_IDLE, 3)  # low
         try:
-            p.ionice(2)
-            ioclass, value = p.ionice()
-            if enum is not None:
-                self.assertIsInstance(ioclass, enum.IntEnum)
-            self.assertEqual(ioclass, 2)
-            self.assertEqual(value, 4)
-            #
-            p.ionice(3)
-            ioclass, value = p.ionice()
-            self.assertEqual(ioclass, 3)
-            self.assertEqual(value, 0)
-            #
-            p.ionice(2, 0)
-            ioclass, value = p.ionice()
-            self.assertEqual(ioclass, 2)
-            self.assertEqual(value, 0)
-            p.ionice(2, 7)
-            ioclass, value = p.ionice()
-            self.assertEqual(ioclass, 2)
-            self.assertEqual(value, 7)
+            # low
+            p.ionice(psutil.IOPRIO_CLASS_IDLE)
+            self.assertEqual(tuple(p.ionice()), (psutil.IOPRIO_CLASS_IDLE, 0))
+            with self.assertRaises(ValueError):  # accepts no value
+                p.ionice(psutil.IOPRIO_CLASS_IDLE, value=7)
+            # normal
+            p.ionice(psutil.IOPRIO_CLASS_BE)
+            self.assertEqual(tuple(p.ionice()), (psutil.IOPRIO_CLASS_BE, 0))
+            p.ionice(psutil.IOPRIO_CLASS_BE, value=7)
+            self.assertEqual(tuple(p.ionice()), (psutil.IOPRIO_CLASS_BE, 7))
+            with self.assertRaises(ValueError):
+                p.ionice(psutil.IOPRIO_CLASS_BE, value=8)
+            # high
+            if os.getuid() == 0:  # root
+                p.ionice(psutil.IOPRIO_CLASS_RT)
+                self.assertEqual(tuple(p.ionice()),
+                                 (psutil.IOPRIO_CLASS_RT, 0))
+                p.ionice(psutil.IOPRIO_CLASS_RT, value=7)
+                self.assertEqual(tuple(p.ionice()),
+                                 (psutil.IOPRIO_CLASS_RT, 7))
+                with self.assertRaises(ValueError):
+                    p.ionice(psutil.IOPRIO_CLASS_IDLE, value=8)
         finally:
-            p.ionice(psutil.IOPRIO_CLASS_NONE, value=0)
+            p.ionice(psutil.IOPRIO_CLASS_BE)
 
     @unittest.skipIf(not HAS_IONICE, "not supported")
     @unittest.skipIf(not WINDOWS, 'not supported on this win version')