PYTHON-3390 Test for encrypted client post-fork (#1037)

Author: bcwarner

test/test_encryption.py

@@ -329,6 +329,25 @@ def test_use_after_close(self):
         with self.assertRaisesRegex(InvalidOperation, "Cannot use MongoClient after close"):
             client.admin.command("ping")
 
+    # Not available for versions of Python without "register_at_fork"
+    @unittest.skipIf(
+        not hasattr(os, "register_at_fork"),
+        "register_at_fork not available in this version of Python",
+    )
+    def test_fork(self):
+        opts = AutoEncryptionOpts(KMS_PROVIDERS, "keyvault.datakeys")
+        client = rs_or_single_client(auto_encryption_opts=opts)
+
+        lock_pid = os.fork()
+        if lock_pid == 0:
+            client.admin.command("ping")
+            client.close()
+            os._exit(0)
+        else:
+            self.assertEqual(0, os.waitpid(lock_pid, 0)[1])
+            client.admin.command("ping")
+            client.close()
+
 
 class TestEncryptedBulkWrite(BulkTestBase, EncryptionIntegrationTest):
     def test_upsert_uuid_standard_encrypt(self):

test/test_fork.py

@@ -53,7 +53,7 @@ def exit_cond():
             if lock_pid == 0:
                 os._exit(exit_cond())
             else:
-                self.assertEqual(0, os.waitpid(lock_pid, 0)[1] >> 8)
+                self.assertEqual(0, os.waitpid(lock_pid, 0)[1])
 
     def test_lock_object_id(self):
         """
@@ -67,7 +67,7 @@ def test_lock_object_id(self):
             if lock_pid == 0:
                 os._exit(int(ObjectId._inc_lock.locked()))
             else:
-                self.assertEqual(0, os.waitpid(lock_pid, 0)[1] >> 8)
+                self.assertEqual(0, os.waitpid(lock_pid, 0)[1])
 
     def test_topology_reset(self):
         """
@@ -92,7 +92,7 @@ def test_topology_reset(self):
             )
             os._exit(0)
         else:  # Parent
-            self.assertEqual(0, os.waitpid(lock_pid, 0)[1] >> 8)
+            self.assertEqual(0, os.waitpid(lock_pid, 0)[1])
             self.assertEqual(self.client._topology._pid, init_id)
             child_id = parent_conn.recv()
             self.assertNotEqual(child_id, init_id)