pylib/util.py

@@ -13,6 +13,15 @@ import socket
 import sys
 import time
 
+if "get_terminal_size" not in dir(shutil):
+    # used by termsize() on python 2.x systems
+    import fcntl
+    import termios
+    import struct
+    PY3 = False
+else:
+    PY3 = True
+
 import ntp.ntpc
 import ntp.version
 import ntp.magic
@@ -532,15 +541,12 @@ def termsize():
     # The way this is used makes it not a big deal if the default is wrong.
     size = (80, 24)
     if os.isatty(1):
-        try:
+        if PY3 is True:
             (w, h) = shutil.get_terminal_size((80, 24))
             size = (w, h)
-        except AttributeError:
+        else:
             try:
                 # OK, Python version < 3.3, cope
-                import fcntl
-                import termios
-                import struct
                 h, w, hp, wp = struct.unpack(
                     'HHHH',
                     fcntl.ioctl(2, termios.TIOCGWINSZ,

tests/pylib/jigs.py

@@ -190,3 +190,34 @@ class SelectModuleJig:
             return (ins, [], [])
         else:
             return ([], [], [])
+
+
+class OSModuleJig:
+    def __init__(self):
+        self.isatty_calls = []
+        self.isatty_returns = []
+
+    def isatty(self, fd):
+        self.isatty_calls.append(fd)
+        return self.isatty_returns.pop(0)
+
+
+class FcntlModuleJig:
+    def __init__(self):
+        self.ioctl_calls = []
+        self.ioctl_returns = []
+
+    def ioctl(self, fd, op, arg=0, mutate_flag=False):
+        self.ioctl_calls.append((fd, op, arg, mutate_flag))
+        return self.ioctl_returns.pop(0)
+
+
+class ShutilModuleJig:
+    def __init__(self):
+        self.gts_calls = []
+        self.gts_returns = []
+
+    def get_terminal_size(self, default=(80, 24)):
+        self.gts_calls.append(default)
+        return self.gts_returns.pop(0)
+    

tests/pylib/test_util.py

@@ -3,6 +3,8 @@
 
 import unittest
 import ntp.util
+import shutil
+import sys
 
 import jigs
 
@@ -37,14 +39,19 @@ class TestPylibUtilMethods(unittest.TestCase):
         self.assertEqual((jig.written, jig.flushed), ("blah", True))
 
     def test_rfc3339(self):
-        self.assertEqual(ntp.util.rfc3339(1480999786),
+        f = ntp.util.rfc3339
+
+        self.assertEqual(f(1480999786),
                          '2016-12-06T04:49:46Z')
-        self.assertEqual(ntp.util.rfc3339(1480999786.5),
+        self.assertEqual(f(1480999786.5),
                          '2016-12-06T04:49:46.5Z')
         # RFC 3339, 2 digits of seconds.
         # we round, but the spec is silent on rounding
-        self.assertEqual(ntp.util.rfc3339(1480999786.025),
-                         '2016-12-06T04:49:46.03Z')
+        # Python 2 and 3 round differently
+        if sys.version_info[0] < 3:
+            self.assertEqual(f(1480999786.025), "2016-12-06T04:49:46.03Z")
+        else:
+            self.assertEqual(f(1480999786.025), "2016-12-06T04:49:46.025Z")
 
     def test_slicedata(self):
         f = ntp.util.slicedata
@@ -455,7 +462,45 @@ class TestPylibUtilMethods(unittest.TestCase):
 
     def test_termsize(self):
         f = ntp.util.termsize
-        # TODO: write this, it needs many jigs
+
+        fakeosmod = jigs.OSModuleJig()
+        fakefcntlmod = jigs.FcntlModuleJig()
+        fakeshutilmod = jigs.ShutilModuleJig()
+        try:
+            ostemp = ntp.util.os
+            ntp.util.os = fakeosmod
+            # Test default
+            fakeosmod.isatty_returns = [False]
+            self.assertEqual(f(), (80, 24))
+            self.assertEqual(fakeosmod.isatty_calls, [1])
+            # termsize takes different code paths for different
+            # versions of Python
+            if "get_terminal_size" in dir(shutil):
+                # Python 3 version
+                try:
+                    shutiltemp = ntp.util.shutil
+                    ntp.util.shutil = fakeshutilmod
+                    fakeosmod.isatty_returns = [True]
+                    fakeshutilmod.gts_returns = [(42, 23)]
+                    self.assertEqual(f(), (42, 23))
+                finally:
+                    ntp.util.shutil = shutiltemp
+            else:
+                # Python 2.x version
+                try:
+                    fcntltemp = ntp.util.fcntl
+                    ntp.util.fcntl = fakefcntlmod
+                    fakeosmod.isatty_returns = [True]
+                    data = ["\x11\x11\x22\x22\x33\x33\x44\x44"]
+                    fakefcntlmod.ioctl_returns = data
+                    self.assertEqual(f(), (0x2222, 0x1111))
+                    self.assertEqual(fakefcntlmod.ioctl_calls,
+                                     [(2, ntp.util.termios.TIOCGWINSZ,
+                                       "\0\0\0\0\0\0\0\0", False)])
+                finally:
+                    ntp.util.fcntl = fcntltemp
+        finally:
+            ntp.util.os = ostemp
 
     def test_PeerStatusWord(self):
         c = ntp.util.PeerStatusWord