Get big sftp tests passing w/ the sftp client + folder crap being a fixture

3 files changed, 131 insertions, 114 deletions

diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 00000000..dbf2cb0f
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,64 @@
+import os
+import threading
+
+import pytest
+from paramiko import RSAKey, SFTPServer, SFTP, Transport
+
+from .loop import LoopSocket
+from .stub_sftp import StubServer, StubSFTPServer
+from .util import _support
+
+
+# TODO: not a huge fan of conftest.py files, see if we can move these somewhere
+# 'nicer'.
+
+
+def make_sftp_folder(client):
+    """
+    Create some non-existing, new folder on the given SFTP connection.
+    """
+    path = os.environ.get('TEST_FOLDER', 'temp-testing000')
+    # TODO: this is disgusting and old, replace with something smarter/simpler
+    for i in range(1000):
+        path = path[:-3] + '%03d' % i
+        try:
+            client.mkdir(path)
+            return path
+        except (IOError, OSError):
+            pass
+
+
+# TODO: apply at module or session level
+# TODO: roll in SFTP folder setup and teardown?
+# NOTE: This is defined here for use by both SFTP (normal & 'big') suites.
+@pytest.fixture
+def sftp():
+    """
+    Set up an in-memory SFTP server, returning its corresponding SFTPClient.
+    """
+    # Sockets & transports
+    socks = LoopSocket()
+    sockc = LoopSocket()
+    sockc.link(socks)
+    tc = Transport(sockc)
+    ts = Transport(socks)
+    # Auth
+    host_key = RSAKey.from_private_key_file(_support('test_rsa.key'))
+    ts.add_server_key(host_key)
+    # Server & client setup
+    event = threading.Event()
+    server = StubServer()
+    ts.set_subsystem_handler('sftp', SFTPServer, StubSFTPServer)
+    ts.start_server(event, server)
+    tc.connect(username='slowdive', password='pygmalion')
+    event.wait(1.0)
+    client = SFTP.from_transport(tc)
+    # Work in 'remote' folder setup (as it wants to use the client)
+    # TODO: how cleanest to make this available to tests? Doing it this way is
+    # marginally less bad than the previous 'global'-using setup, but not by
+    # much?
+    client.FOLDER = make_sftp_folder(client)
+    # Yield client to caller
+    yield client
+    # Clean up
+    client.rmdir(client.FOLDER)
diff --git a/tests/test_sftp.py b/tests/test_sftp.py
index 432d3e40..a38c2e7d 100755
--- a/tests/test_sftp.py
+++ b/tests/test_sftp.py
@@ -32,6 +32,8 @@ import warnings
 from binascii import hexlify
 from tempfile import mkstemp
 
+import pytest
+
 import paramiko
 import paramiko.util
 from paramiko.py3compat import PY2, b, u, StringIO
@@ -40,7 +42,6 @@ from paramiko.sftp_attr import SFTPAttributes
 
 from .util import needs_builtin
 from .stub_sftp import StubServer, StubSFTPServer
-from .loop import LoopSocket
 from .util import _support
 
 
@@ -94,37 +95,7 @@ unicode_folder = eval(compile(r"u'\u00fcnic\u00f8de'" if PY2 else r"'\u00fcnic\u
 utf8_folder = b'/\xc3\xbcnic\xc3\xb8\x64\x65'
 
 
-# TODO: turn into a pytest fixture; consider making it module or session-global
-# to mimic old behavior (though that still feels unclean to me...)
-def make_loopback_sftp():
-    """
-    Set up an in-memory SFTP server.
-
-    :returns:
-        A 2-tuple of the resulting SFTPClient (for tests that just care about a
-        'default' client) and Transport (for testing instantiation _of_
-        SFTPClient itself, which can take an existing transport object.)
-    """
-    socks = LoopSocket()
-    sockc = LoopSocket()
-    sockc.link(socks)
-    tc = paramiko.Transport(sockc)
-    ts = paramiko.Transport(socks)
-
-    host_key = paramiko.RSAKey.from_private_key_file(_support('test_rsa.key'))
-    ts.add_server_key(host_key)
-    event = threading.Event()
-    server = StubServer()
-    ts.set_subsystem_handler('sftp', paramiko.SFTPServer, StubSFTPServer)
-    ts.start_server(event, server)
-    tc.connect(username='slowdive', password='pygmalion')
-    event.wait(1.0)
-
-    return paramiko.SFTP.from_transport(tc), tc
-
-
-class SFTPTest (unittest.TestCase):
-
+class SFTPTest(unittest.TestCase):
     @staticmethod
     def set_big_file_test(onoff):
         global g_big_file_test
diff --git a/tests/test_sftp_big.py b/tests/test_sftp_big.py
index ef12b05c..c58a7912 100644
--- a/tests/test_sftp_big.py
+++ b/tests/test_sftp_big.py
@@ -32,92 +32,74 @@ import unittest
 
 from paramiko.common import o660
 
-from .test_sftp import make_loopback_sftp
-
 
 FOLDER = os.environ.get('TEST_FOLDER', 'temp-testing000')
 
 
-class BigSFTPTest (unittest.TestCase):
-
-    def setUp(self):
-        global FOLDER
-        self.sftp, _ = make_loopback_sftp()
-        # TODO: same TODOs as in test_sftp.py re: not doing this awful crap
-        for i in range(1000):
-            FOLDER = FOLDER[:-3] + '%03d' % i
-            try:
-                self.sftp.mkdir(FOLDER)
-                break
-            except (IOError, OSError):
-                pass
-
-    def tearDown(self):
-        self.sftp.rmdir(FOLDER)
-
-    def test_1_lots_of_files(self):
+class TestBigSFTP(object):
+    def test_1_lots_of_files(self, sftp):
         """
         create a bunch of files over the same session.
         """
         numfiles = 100
         try:
             for i in range(numfiles):
-                with self.sftp.open('%s/file%d.txt' % (FOLDER, i), 'w', 1) as f:
+                with sftp.open('%s/file%d.txt' % (FOLDER, i), 'w', 1) as f:
                     f.write('this is file #%d.\n' % i)
-                self.sftp.chmod('%s/file%d.txt' % (FOLDER, i), o660)
+                sftp.chmod('%s/file%d.txt' % (FOLDER, i), o660)
 
             # now make sure every file is there, by creating a list of filenmes
             # and reading them in random order.
             numlist = list(range(numfiles))
             while len(numlist) > 0:
                 r = numlist[random.randint(0, len(numlist) - 1)]
-                with self.sftp.open('%s/file%d.txt' % (FOLDER, r)) as f:
-                    self.assertEqual(f.readline(), 'this is file #%d.\n' % r)
+                with sftp.open('%s/file%d.txt' % (FOLDER, r)) as f:
+                    assert f.readline() == 'this is file #%d.\n' % r
                 numlist.remove(r)
         finally:
             for i in range(numfiles):
                 try:
-                    self.sftp.remove('%s/file%d.txt' % (FOLDER, i))
+                    sftp.remove('%s/file%d.txt' % (FOLDER, i))
                 except:
                     pass
 
-    def test_2_big_file(self):
+    def test_2_big_file(self, sftp):
         """
         write a 1MB file with no buffering.
         """
         kblob = (1024 * b'x')
         start = time.time()
         try:
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'w') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'w') as f:
                 for n in range(1024):
                     f.write(kblob)
                     if n % 128 == 0:
                         sys.stderr.write('.')
             sys.stderr.write(' ')
 
-            self.assertEqual(self.sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
+            assert sftp.stat('%s/hongry.txt' % FOLDER).st_size == 1024 * 1024
             end = time.time()
             sys.stderr.write('%ds ' % round(end - start))
             
             start = time.time()
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'r') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'r') as f:
                 for n in range(1024):
                     data = f.read(1024)
-                    self.assertEqual(data, kblob)
+                    assert data == kblob
 
             end = time.time()
             sys.stderr.write('%ds ' % round(end - start))
         finally:
-            self.sftp.remove('%s/hongry.txt' % FOLDER)
+            sftp.remove('%s/hongry.txt' % FOLDER)
 
-    def test_3_big_file_pipelined(self):
+    def test_3_big_file_pipelined(self, sftp):
         """
         write a 1MB file, with no linefeeds, using pipelining.
         """
         kblob = bytes().join([struct.pack('>H', n) for n in range(512)])
         start = time.time()
         try:
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
                 f.set_pipelined(True)
                 for n in range(1024):
                     f.write(kblob)
@@ -125,12 +107,12 @@ class BigSFTPTest (unittest.TestCase):
                         sys.stderr.write('.')
             sys.stderr.write(' ')
 
-            self.assertEqual(self.sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
+            assert sftp.stat('%s/hongry.txt' % FOLDER).st_size == 1024 * 1024
             end = time.time()
             sys.stderr.write('%ds ' % round(end - start))
             
             start = time.time()
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
                 file_size = f.stat().st_size
                 f.prefetch(file_size)
 
@@ -144,18 +126,18 @@ class BigSFTPTest (unittest.TestCase):
                         chunk = size - n
                     data = f.read(chunk)
                     offset = n % 1024
-                    self.assertEqual(data, k2blob[offset:offset + chunk])
+                    assert data == k2blob[offset:offset + chunk]
                     n += chunk
 
             end = time.time()
             sys.stderr.write('%ds ' % round(end - start))
         finally:
-            self.sftp.remove('%s/hongry.txt' % FOLDER)
+            sftp.remove('%s/hongry.txt' % FOLDER)
 
-    def test_4_prefetch_seek(self):
+    def test_4_prefetch_seek(self, sftp):
         kblob = bytes().join([struct.pack('>H', n) for n in range(512)])
         try:
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
                 f.set_pipelined(True)
                 for n in range(1024):
                     f.write(kblob)
@@ -163,13 +145,13 @@ class BigSFTPTest (unittest.TestCase):
                         sys.stderr.write('.')
             sys.stderr.write(' ')
             
-            self.assertEqual(self.sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
+            assert sftp.stat('%s/hongry.txt' % FOLDER).st_size == 1024 * 1024
             
             start = time.time()
             k2blob = kblob + kblob
             chunk = 793
             for i in range(10):
-                with self.sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
+                with sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
                     file_size = f.stat().st_size
                     f.prefetch(file_size)
                     base_offset = (512 * 1024) + 17 * random.randint(1000, 2000)
@@ -181,17 +163,17 @@ class BigSFTPTest (unittest.TestCase):
                         f.seek(offset)
                         data = f.read(chunk)
                         n_offset = offset % 1024
-                        self.assertEqual(data, k2blob[n_offset:n_offset + chunk])
+                        assert data == k2blob[n_offset:n_offset + chunk]
                         offset += chunk
             end = time.time()
             sys.stderr.write('%ds ' % round(end - start))
         finally:
-            self.sftp.remove('%s/hongry.txt' % FOLDER)
+            sftp.remove('%s/hongry.txt' % FOLDER)
 
-    def test_5_readv_seek(self):
+    def test_5_readv_seek(self, sftp):
         kblob = bytes().join([struct.pack('>H', n) for n in range(512)])
         try:
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
                 f.set_pipelined(True)
                 for n in range(1024):
                     f.write(kblob)
@@ -199,13 +181,13 @@ class BigSFTPTest (unittest.TestCase):
                         sys.stderr.write('.')
             sys.stderr.write(' ')
 
-            self.assertEqual(self.sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
+            assert sftp.stat('%s/hongry.txt' % FOLDER).st_size == 1024 * 1024
 
             start = time.time()
             k2blob = kblob + kblob
             chunk = 793
             for i in range(10):
-                with self.sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
+                with sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
                     base_offset = (512 * 1024) + 17 * random.randint(1000, 2000)
                     # make a bunch of offsets and put them in random order
                     offsets = [base_offset + j * chunk for j in range(100)]
@@ -218,20 +200,20 @@ class BigSFTPTest (unittest.TestCase):
                     for i in range(len(readv_list)):
                         offset = readv_list[i][0]
                         n_offset = offset % 1024
-                        self.assertEqual(next(ret), k2blob[n_offset:n_offset + chunk])
+                        assert next(ret) == k2blob[n_offset:n_offset + chunk]
             end = time.time()
             sys.stderr.write('%ds ' % round(end - start))
         finally:
-            self.sftp.remove('%s/hongry.txt' % FOLDER)
+            sftp.remove('%s/hongry.txt' % FOLDER)
 
-    def test_6_lots_of_prefetching(self):
+    def test_6_lots_of_prefetching(self, sftp):
         """
         prefetch a 1MB file a bunch of times, discarding the file object
         without using it, to verify that paramiko doesn't get confused.
         """
         kblob = (1024 * b'x')
         try:
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'w') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'w') as f:
                 f.set_pipelined(True)
                 for n in range(1024):
                     f.write(kblob)
@@ -239,31 +221,31 @@ class BigSFTPTest (unittest.TestCase):
                         sys.stderr.write('.')
             sys.stderr.write(' ')
 
-            self.assertEqual(self.sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
+            assert sftp.stat('%s/hongry.txt' % FOLDER).st_size == 1024 * 1024
 
             for i in range(10):
-                with self.sftp.open('%s/hongry.txt' % FOLDER, 'r') as f:
+                with sftp.open('%s/hongry.txt' % FOLDER, 'r') as f:
                     file_size = f.stat().st_size
                     f.prefetch(file_size)
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'r') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'r') as f:
                 file_size = f.stat().st_size
                 f.prefetch(file_size)
                 for n in range(1024):
                     data = f.read(1024)
-                    self.assertEqual(data, kblob)
+                    assert data == kblob
                     if n % 128 == 0:
                         sys.stderr.write('.')
             sys.stderr.write(' ')
         finally:
-            self.sftp.remove('%s/hongry.txt' % FOLDER)
+            sftp.remove('%s/hongry.txt' % FOLDER)
     
-    def test_7_prefetch_readv(self):
+    def test_7_prefetch_readv(self, sftp):
         """
         verify that prefetch and readv don't conflict with each other.
         """
         kblob = bytes().join([struct.pack('>H', n) for n in range(512)])
         try:
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
                 f.set_pipelined(True)
                 for n in range(1024):
                     f.write(kblob)
@@ -271,13 +253,13 @@ class BigSFTPTest (unittest.TestCase):
                         sys.stderr.write('.')
             sys.stderr.write(' ')
             
-            self.assertEqual(self.sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
+            assert sftp.stat('%s/hongry.txt' % FOLDER).st_size == 1024 * 1024
 
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
                 file_size = f.stat().st_size
                 f.prefetch(file_size)
                 data = f.read(1024)
-                self.assertEqual(data, kblob)
+                assert data == kblob
 
                 chunk_size = 793
                 base_offset = 512 * 1024
@@ -285,22 +267,22 @@ class BigSFTPTest (unittest.TestCase):
                 chunks = [(base_offset + (chunk_size * i), chunk_size) for i in range(20)]
                 for data in f.readv(chunks):
                     offset = base_offset % 1024
-                    self.assertEqual(chunk_size, len(data))
-                    self.assertEqual(k2blob[offset:offset + chunk_size], data)
+                    assert chunk_size == len(data)
+                    assert k2blob[offset:offset + chunk_size] == data
                     base_offset += chunk_size
 
             sys.stderr.write(' ')
         finally:
-            self.sftp.remove('%s/hongry.txt' % FOLDER)
+            sftp.remove('%s/hongry.txt' % FOLDER)
     
-    def test_8_large_readv(self):
+    def test_8_large_readv(self, sftp):
         """
         verify that a very large readv is broken up correctly and still
         returned as a single blob.
         """
         kblob = bytes().join([struct.pack('>H', n) for n in range(512)])
         try:
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
                 f.set_pipelined(True)
                 for n in range(1024):
                     f.write(kblob)
@@ -308,53 +290,53 @@ class BigSFTPTest (unittest.TestCase):
                         sys.stderr.write('.')
             sys.stderr.write(' ')
 
-            self.assertEqual(self.sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
+            assert sftp.stat('%s/hongry.txt' % FOLDER).st_size == 1024 * 1024
             
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
                 data = list(f.readv([(23 * 1024, 128 * 1024)]))
-                self.assertEqual(1, len(data))
+                assert len(data) == 1
                 data = data[0]
-                self.assertEqual(128 * 1024, len(data))
+                assert len(data) == 128 * 1024
             
             sys.stderr.write(' ')
         finally:
-            self.sftp.remove('%s/hongry.txt' % FOLDER)
+            sftp.remove('%s/hongry.txt' % FOLDER)
     
-    def test_9_big_file_big_buffer(self):
+    def test_9_big_file_big_buffer(self, sftp):
         """
         write a 1MB file, with no linefeeds, and a big buffer.
         """
         mblob = (1024 * 1024 * 'x')
         try:
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'w', 128 * 1024) as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'w', 128 * 1024) as f:
                 f.write(mblob)
 
-            self.assertEqual(self.sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
+            assert sftp.stat('%s/hongry.txt' % FOLDER).st_size == 1024 * 1024
         finally:
-            self.sftp.remove('%s/hongry.txt' % FOLDER)
+            sftp.remove('%s/hongry.txt' % FOLDER)
     
-    def test_A_big_file_renegotiate(self):
+    def test_A_big_file_renegotiate(self, sftp):
         """
         write a 1MB file, forcing key renegotiation in the middle.
         """
-        t = self.sftp.sock.get_transport()
+        t = sftp.sock.get_transport()
         t.packetizer.REKEY_BYTES = 512 * 1024
         k32blob = (32 * 1024 * 'x')
         try:
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'w', 128 * 1024) as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'w', 128 * 1024) as f:
                 for i in range(32):
                     f.write(k32blob)
 
-            self.assertEqual(self.sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
-            self.assertNotEqual(t.H, t.session_id)
+            assert sftp.stat('%s/hongry.txt' % FOLDER).st_size == 1024 * 1024
+            assert t.H != t.session_id
             
             # try to read it too.
-            with self.sftp.open('%s/hongry.txt' % FOLDER, 'r', 128 * 1024) as f:
+            with sftp.open('%s/hongry.txt' % FOLDER, 'r', 128 * 1024) as f:
                 file_size = f.stat().st_size
                 f.prefetch(file_size)
                 total = 0
                 while total < 1024 * 1024:
                     total += len(f.read(32 * 1024))
         finally:
-            self.sftp.remove('%s/hongry.txt' % FOLDER)
+            sftp.remove('%s/hongry.txt' % FOLDER)
             t.packetizer.REKEY_BYTES = pow(2, 30)