From 531e404879f0f077d0f1a0d9231740ac58eaa165 Mon Sep 17 00:00:00 2001
From: Dave Brondsema <dave@brondsema.net>
Date: Thu, 2 Jun 2022 12:11:06 -0400
Subject: [PATCH] test improvements - from #45

---
 ming/tests/odm/test_declarative.py | 33 ++++++++++++++++++++++++++++++
 ming/tests/odm/test_mapper.py      | 17 +++++++++++++++
 ming/tests/test_session.py         |  2 +-
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/ming/tests/odm/test_declarative.py b/ming/tests/odm/test_declarative.py
index 876ae75..e98f315 100644
--- a/ming/tests/odm/test_declarative.py
+++ b/ming/tests/odm/test_declarative.py
@@ -837,3 +837,36 @@ def test_hook_base(self):
         Mapper.replace_session(new_session)
         assert id(self.Basic.query.session) == id(new_session)
         assert id(self.session) != id(new_session)
+
+
+class TestBeforeSave(TestCase):
+
+    def setUp(self):
+        Mapper._mapper_by_classname.clear()
+        self.datastore = create_datastore('mim:///test_db')
+        self.session = ODMSession(bind=self.datastore)
+        class Basic(MappedClass):
+            class __mongometa__:
+                name = 'hook'
+                session = self.session
+                def before_save(instance):
+                    instance.a = 9
+
+            _id = FieldProperty(S.ObjectId)
+            a = FieldProperty(int)
+        Mapper.compile_all()
+        self.Basic = Basic
+        self.session.remove(self.Basic)
+
+    def test_hook_base(self):
+        doc = self.Basic()
+        doc.a = 5
+        self.session.flush()  # first insert
+        self.session.close()
+        doc = self.Basic.query.get(doc._id)
+        assert doc.a == 9, doc.a
+        doc.a = 6
+        self.session.flush()  # then save
+        self.session.close()
+        doc = self.Basic.query.get(doc._id)
+        assert doc.a == 9, doc.a
\ No newline at end of file
diff --git a/ming/tests/odm/test_mapper.py b/ming/tests/odm/test_mapper.py
index e8055a7..0c202da 100644
--- a/ming/tests/odm/test_mapper.py
+++ b/ming/tests/odm/test_mapper.py
@@ -277,6 +277,23 @@ def test_group(self, pymongo_group):
         self.Basic.query.group()
         assert pymongo_group.called
 
+    def test_multiple_update_flushes(self):
+        initial_doc = self.Basic()
+        initial_doc.a = 1
+        self.session.flush()
+        self.session.close()
+
+        doc_updating = self.Basic.query.get(_id=initial_doc._id)
+        doc_updating.a = 2
+        self.session.flush()
+        doc_updating.a = 1  # back to "initial" value
+        doc_updating.e = 'foo'  # change something else too
+        self.session.flush()
+        self.session.close()
+
+        doc_after_updates = self.Basic.query.get(_id=doc_updating._id)
+        assert doc_after_updates.a == 1
+
 
 class TestRelation(TestCase):
     def setUp(self):
diff --git a/ming/tests/test_session.py b/ming/tests/test_session.py
index 4631e25..80fe412 100644
--- a/ming/tests/test_session.py
+++ b/ming/tests/test_session.py
@@ -81,7 +81,7 @@ def test_base_session(self):
         doc = self.TestDocNoSchema({'_id':5, 'a':5})
         sess.save(doc)
         impl.save.assert_called_with(dict(_id=5, a=5))
-        doc = self.TestDocNoSchema({'_id':5, 'a':5})
+        doc = self.TestDocNoSchema({'_id':5, 'a':5, 'b': 6})
         sess.save(doc, 'a')
         impl.update.assert_called_with(dict(_id=5), {'$set':dict(a=5)})
         doc = self.TestDocNoSchema({'_id':5, 'a':5})