From 4c62a060f0c460b1d68ebba861eb7a3fd9f9e14c Mon Sep 17 00:00:00 2001
From: Agustin Barto <abarto@gmail.com>
Date: Tue, 17 Mar 2020 19:37:21 -0300
Subject: [PATCH] Add tests for set_write_conern and set_read_write_concern

---
 tests/test_context_managers.py | 48 ++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/tests/test_context_managers.py b/tests/test_context_managers.py
index fa3f5960..f445cf57 100644
--- a/tests/test_context_managers.py
+++ b/tests/test_context_managers.py
@@ -10,11 +10,59 @@ from mongoengine.context_managers import (
     query_counter,
     switch_collection,
     switch_db,
+    set_write_concern,
+    set_read_write_concern,
 )
 from mongoengine.pymongo_support import count_documents
 
 
 class TestContextManagers:
+    def test_set_write_concern(self):
+        connect("mongoenginetest")
+
+        class User(Document):
+            name = StringField()
+
+        collection = User._get_collection()
+        original_write_concern = collection.write_concern
+
+        with set_write_concern(
+            collection, {"w": "majority", "j": True, "wtimeout": 1234}
+        ) as updated_collection:
+            assert updated_collection.write_concern.document == {
+                "w": "majority",
+                "j": True,
+                "wtimeout": 1234,
+            }
+
+        assert original_write_concern.document == collection.write_concern.document
+
+    def test_set_read_write_concern(self):
+        connect("mongoenginetest")
+
+        class User(Document):
+            name = StringField()
+
+        collection = User._get_collection()
+
+        original_read_concern = collection.read_concern
+        original_write_concern = collection.write_concern
+
+        with set_read_write_concern(
+            collection,
+            {"w": "majority", "j": True, "wtimeout": 1234},
+            {"level": "local"},
+        ) as update_collection:
+            assert update_collection.read_concern.document == {"level": "local"}
+            assert update_collection.write_concern.document == {
+                "w": "majority",
+                "j": True,
+                "wtimeout": 1234,
+            }
+
+        assert original_read_concern.document == collection.read_concern.document
+        assert original_write_concern.document == collection.write_concern.document
+
     def test_switch_db_context_manager(self):
         connect("mongoenginetest")
         register_connection("testdb-1", "mongoenginetest2")