Search code examples
pythonunit-testingsshpython-unittest

Python unit tests run function after all test


I need to test something on Python via ssh. I don't want to make ssh connection for every test, because it is too long, I have written this:

class TestCase(unittest.TestCase):
    client = None
    def setUp(self):
        if not hasattr(self.__class__, 'client') or self.__class__.client is None:
            self.__class__.client = paramiko.SSHClient()
            self.__class__.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            self.__class__.client.connect(hostname=consts.get_host(), port=consts.get_port(), username=consts.get_user(),
                                password=consts.get_password())

    def test_a(self):
        pass

    def test_b(self):
        pass

    def test_c(self):
        pass

    def disconnect(self):
        self.__class__.client.close()

and my runner

if __name__ == '__main__':
    suite = unittest.TestSuite((
        unittest.makeSuite(TestCase),
    ))
    result = unittest.TextTestRunner().run(suite)
    TestCase.disconnect()
    sys.exit(not result.wasSuccessful())

In this version I get error TypeError: unbound method disconnect() must be called with TestCase instance as first argument (got nothing instead). So how can I call disconnect after all tests pass?


Solution

  • You should use setUpClass and tearDownClass instead, if you want to keep the same connection for all tests. You'll also need to make the disconnect method static, so it belongs to the class and not an instance of the class.

    class TestCase(unittest.TestCase):
    
         def setUpClass(cls):
             cls.connection = <your connection setup>
    
         @staticmethod
         def disconnect():
             ... disconnect TestCase.connection
    
         def tearDownClass(cls):
             cls.disconnect()