gh-106527: asyncio: optimize to add/remove readers and writers (#106528)

1 files changed, 18 insertions, 18 deletions

diff --git a/Lib/test/test_asyncio/test_selector_events.py b/Lib/test/test_asyncio/test_selector_events.py
index 47693ea..c22b780 100644
--- a/Lib/test/test_asyncio/test_selector_events.py
+++ b/Lib/test/test_asyncio/test_selector_events.py
@@ -178,7 +178,7 @@ class BaseSelectorEventLoopTests(test_utils.TestCase):
                 sock.connect.assert_called_with(('127.0.0.1', 0))
 
     def test_add_reader(self):
-        self.loop._selector.get_key.side_effect = KeyError
+        self.loop._selector.get_map.return_value = {}
         cb = lambda: True
         self.loop.add_reader(1, cb)
 
@@ -192,8 +192,8 @@ class BaseSelectorEventLoopTests(test_utils.TestCase):
     def test_add_reader_existing(self):
         reader = mock.Mock()
         writer = mock.Mock()
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
-            1, 1, selectors.EVENT_WRITE, (reader, writer))
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
+            1, 1, selectors.EVENT_WRITE, (reader, writer))}
         cb = lambda: True
         self.loop.add_reader(1, cb)
 
@@ -208,8 +208,8 @@ class BaseSelectorEventLoopTests(test_utils.TestCase):
 
     def test_add_reader_existing_writer(self):
         writer = mock.Mock()
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
-            1, 1, selectors.EVENT_WRITE, (None, writer))
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
+            1, 1, selectors.EVENT_WRITE, (None, writer))}
         cb = lambda: True
         self.loop.add_reader(1, cb)
 
@@ -222,8 +222,8 @@ class BaseSelectorEventLoopTests(test_utils.TestCase):
         self.assertEqual(writer, w)
 
     def test_remove_reader(self):
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
-            1, 1, selectors.EVENT_READ, (None, None))
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
+            1, 1, selectors.EVENT_READ, (None, None))}
         self.assertFalse(self.loop.remove_reader(1))
 
         self.assertTrue(self.loop._selector.unregister.called)
@@ -231,9 +231,9 @@ class BaseSelectorEventLoopTests(test_utils.TestCase):
     def test_remove_reader_read_write(self):
         reader = mock.Mock()
         writer = mock.Mock()
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
             1, 1, selectors.EVENT_READ | selectors.EVENT_WRITE,
-            (reader, writer))
+            (reader, writer))}
         self.assertTrue(
             self.loop.remove_reader(1))
 
@@ -243,12 +243,12 @@ class BaseSelectorEventLoopTests(test_utils.TestCase):
             self.loop._selector.modify.call_args[0])
 
     def test_remove_reader_unknown(self):
-        self.loop._selector.get_key.side_effect = KeyError
+        self.loop._selector.get_map.return_value = {}
         self.assertFalse(
             self.loop.remove_reader(1))
 
     def test_add_writer(self):
-        self.loop._selector.get_key.side_effect = KeyError
+        self.loop._selector.get_map.return_value = {}
         cb = lambda: True
         self.loop.add_writer(1, cb)
 
@@ -262,8 +262,8 @@ class BaseSelectorEventLoopTests(test_utils.TestCase):
     def test_add_writer_existing(self):
         reader = mock.Mock()
         writer = mock.Mock()
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
-            1, 1, selectors.EVENT_READ, (reader, writer))
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
+            1, 1, selectors.EVENT_READ, (reader, writer))}
         cb = lambda: True
         self.loop.add_writer(1, cb)
 
@@ -277,8 +277,8 @@ class BaseSelectorEventLoopTests(test_utils.TestCase):
         self.assertEqual(cb, w._callback)
 
     def test_remove_writer(self):
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
-            1, 1, selectors.EVENT_WRITE, (None, None))
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
+            1, 1, selectors.EVENT_WRITE, (None, None))}
         self.assertFalse(self.loop.remove_writer(1))
 
         self.assertTrue(self.loop._selector.unregister.called)
@@ -286,9 +286,9 @@ class BaseSelectorEventLoopTests(test_utils.TestCase):
     def test_remove_writer_read_write(self):
         reader = mock.Mock()
         writer = mock.Mock()
-        self.loop._selector.get_key.return_value = selectors.SelectorKey(
+        self.loop._selector.get_map.return_value = {1: selectors.SelectorKey(
             1, 1, selectors.EVENT_READ | selectors.EVENT_WRITE,
-            (reader, writer))
+            (reader, writer))}
         self.assertTrue(
             self.loop.remove_writer(1))
 
@@ -298,7 +298,7 @@ class BaseSelectorEventLoopTests(test_utils.TestCase):
             self.loop._selector.modify.call_args[0])
 
     def test_remove_writer_unknown(self):
-        self.loop._selector.get_key.side_effect = KeyError
+        self.loop._selector.get_map.return_value = {}
         self.assertFalse(
             self.loop.remove_writer(1))