TST: cleanup of pickled cids

tacaswell · tacaswell · commit 3204907a4b43 · 2021-02-18T18:11:39.000-05:00
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -182,8 +182,11 @@ def setup(self):
         self.signal = 'test'
         self.callbacks = cbook.CallbackRegistry()
 
-    def connect(self, s, func):
-        return self.callbacks.connect(s, func)
+    def connect(self, s, func, pickle):
+        cid = self.callbacks.connect(s, func)
+        if pickle:
+            self.callbacks._pickled_cids.add(cid)
+        return cid
 
     def disconnect(self, cid):
         return self.callbacks.disconnect(cid)
@@ -197,25 +200,27 @@ def count(self):
     def is_empty(self):
         assert self.callbacks._func_cid_map == {}
         assert self.callbacks.callbacks == {}
+        assert self.callbacks._pickled_cids == set()
 
     def is_not_empty(self):
         assert self.callbacks._func_cid_map != {}
         assert self.callbacks.callbacks != {}
 
-    def test_callback_complete(self):
+    @pytest.mark.parametrize('pickle', [True, False])
+    def test_callback_complete(self, pickle):
         # ensure we start with an empty registry
         self.is_empty()
 
         # create a class for testing
         mini_me = Test_callback_registry()
 
         # test that we can add a callback
-        cid1 = self.connect(self.signal, mini_me.dummy)
+        cid1 = self.connect(self.signal, mini_me.dummy, pickle)
         assert type(cid1) == int
         self.is_not_empty()
 
         # test that we don't add a second callback
-        cid2 = self.connect(self.signal, mini_me.dummy)
+        cid2 = self.connect(self.signal, mini_me.dummy, pickle)
         assert cid1 == cid2
         self.is_not_empty()
         assert len(self.callbacks._func_cid_map) == 1
@@ -226,15 +231,16 @@ def test_callback_complete(self):
         # check we now have no callbacks registered
         self.is_empty()
 
-    def test_callback_disconnect(self):
+    @pytest.mark.parametrize('pickle', [True, False])
+    def test_callback_disconnect(self, pickle):
         # ensure we start with an empty registry
         self.is_empty()
 
         # create a class for testing
         mini_me = Test_callback_registry()
 
         # test that we can add a callback
-        cid1 = self.connect(self.signal, mini_me.dummy)
+        cid1 = self.connect(self.signal, mini_me.dummy, pickle)
         assert type(cid1) == int
         self.is_not_empty()
 
@@ -243,15 +249,16 @@ def test_callback_disconnect(self):
         # check we now have no callbacks registered
         self.is_empty()
 
-    def test_callback_wrong_disconnect(self):
+    @pytest.mark.parametrize('pickle', [True, False])
+    def test_callback_wrong_disconnect(self, pickle):
         # ensure we start with an empty registry
         self.is_empty()
 
         # create a class for testing
         mini_me = Test_callback_registry()
 
         # test that we can add a callback
-        cid1 = self.connect(self.signal, mini_me.dummy)
+        cid1 = self.connect(self.signal, mini_me.dummy, pickle)
         assert type(cid1) == int
         self.is_not_empty()
 
@@ -260,21 +267,22 @@ def test_callback_wrong_disconnect(self):
         # check we still have callbacks registered
         self.is_not_empty()
 
-    def test_registration_on_non_empty_registry(self):
+    @pytest.mark.parametrize('pickle', [True, False])
+    def test_registration_on_non_empty_registry(self, pickle):
         # ensure we start with an empty registry
         self.is_empty()
 
         # setup the registry with a callback
         mini_me = Test_callback_registry()
-        self.connect(self.signal, mini_me.dummy)
+        self.connect(self.signal, mini_me.dummy, pickle)
 
         # Add another callback
         mini_me2 = Test_callback_registry()
-        self.connect(self.signal, mini_me2.dummy)
+        self.connect(self.signal, mini_me2.dummy, pickle)
 
         # Remove and add the second callback
         mini_me2 = Test_callback_registry()
-        self.connect(self.signal, mini_me2.dummy)
+        self.connect(self.signal, mini_me2.dummy, pickle)
 
         # We still have 2 references
         self.is_not_empty()
