Be more careful with handle cleanup, to support parallel testing.

1 files changed, 9 insertions, 8 deletions

diff --git a/unit_tests.py b/unit_tests.py
index f9bff67..2fc9d3f 100644
--- a/unit_tests.py
+++ b/unit_tests.py
@@ -216,19 +216,22 @@ class TestKeys(TestCase):
         p11.C_Finalize()
 
     def setUp(self):
+        self.handles = set()
         self.session = p11.C_OpenSession(args.slot)
         p11.C_Login(self.session, CKU_USER, args.user_pin)
         super(TestKeys, self).setUp()
 
     def tearDown(self):
         super(TestKeys, self).tearDown()
-        for handle in p11.FindObjects(self.session):
+        for handle in self.handles:
             p11.C_DestroyObject(self.session, handle)
         p11.C_CloseAllSessions(args.slot)
         del self.session
 
     def assertIsKeypair(self, *keypair):
         public_handle, private_handle = keypair[0] if isinstance(keypair[0], tuple) else keypair
+        self.handles.add(public_handle)
+        self.handles.add(private_handle)
         self.assertEqual(p11.C_GetAttributeValue(self.session, public_handle,  CKA_CLASS), {CKA_CLASS: CKO_PUBLIC_KEY})
         self.assertEqual(p11.C_GetAttributeValue(self.session, private_handle, CKA_CLASS), {CKA_CLASS: CKO_PRIVATE_KEY})
 
@@ -245,11 +248,6 @@ class TestKeys(TestCase):
                                 CKA_ID = "EC-P256", CKA_EC_PARAMS = self.oid_p256,
                                 CKA_SIGN = True, CKA_VERIFY = True))
 
-        # Might need to do this until we expand the number of key slots
-        if False:
-          for handle in p11.FindObjects(self.session):
-            p11.C_DestroyObject(self.session, handle)
-
         self.assertIsKeypair(
             p11.C_GenerateKeyPair(self.session, CKM_EC_KEY_PAIR_GEN,
                                   public_CKA_TOKEN = False, private_CKA_TOKEN = True,
@@ -359,6 +357,7 @@ class TestKeys(TestCase):
           CKA_VERIFY        = True,
           CKA_EC_POINT      = Q,
           CKA_EC_PARAMS     = self.oid_p256)
+        self.handles.add(handle)
         p11.C_VerifyInit(self.session, CKM_ECDSA, handle)
         p11.C_Verify(self.session, H, r + s)
 
@@ -378,6 +377,7 @@ class TestKeys(TestCase):
           CKA_VERIFY        = True,
           CKA_EC_POINT      = Q,
           CKA_EC_PARAMS     = self.oid_p384)
+        self.handles.add(handle)
         p11.C_VerifyInit(self.session, CKM_ECDSA, handle)
         p11.C_Verify(self.session, H, r + s)
 
@@ -399,13 +399,14 @@ class TestKeys(TestCase):
                                     CKA_CLASS, CKA_KEY_TYPE, CKA_VERIFY, CKA_TOKEN,
                                     CKA_EC_PARAMS, CKA_EC_POINT)
 
-        for handle in p11.FindObjects(self.session):
-            p11.C_DestroyObject(self.session, handle)
+        while self.handles:
+            p11.C_DestroyObject(self.session, self.handles.pop())
         p11.C_CloseAllSessions(args.slot)
         self.session = p11.C_OpenSession(args.slot)
         p11.C_Login(self.session, CKU_USER, args.user_pin)
 
         o = p11.C_CreateObject(self.session, a)
+        self.handles.add(o)
         p11.C_VerifyInit(self.session, CKM_ECDSA_SHA256, o)
         p11.C_Verify(self.session, hamster, sig)