fix: Improve test cleanup. (#234)

Make sure that `dispose_of` is called immediately after storing an
entity to make sure an error in the test doesn't prevent the entities
from getting cleaned up.

tests/system/test_crud.py

@@ -219,6 +219,7 @@ class SomeKind(ndb.Model):
 
     entity = SomeKind(foo=42, bar="none")
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == 42
@@ -228,8 +229,6 @@ class SomeKind(ndb.Model):
     ds_entity = ds_client.get(key._key)
     assert ds_entity["bar"] == "none"
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_insert_entity_with_stored_name_property(dispose_of, ds_client):
@@ -239,6 +238,7 @@ class SomeKind(ndb.Model):
 
     entity = SomeKind(foo="something", bar="or other")
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == "something"
@@ -247,8 +247,6 @@ class SomeKind(ndb.Model):
     ds_entity = ds_client.get(key._key)
     assert ds_entity["notbar"] == "or other"
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_insert_roundtrip_naive_datetime(dispose_of, ds_client):
@@ -257,12 +255,11 @@ class SomeKind(ndb.Model):
 
     entity = SomeKind(foo=datetime.datetime(2010, 5, 12, 2, 42))
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == datetime.datetime(2010, 5, 12, 2, 42)
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_datetime_w_tzinfo(dispose_of, ds_client):
@@ -287,13 +284,12 @@ class SomeKind(ndb.Model):
         bar=datetime.datetime(2010, 5, 12, 2, 42),
     )
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == datetime.datetime(2010, 5, 12, 3, 42, tzinfo=mytz)
     assert retrieved.bar == datetime.datetime(2010, 5, 11, 22, 42, tzinfo=mytz)
 
-    dispose_of(key._key)
-
 
 def test_parallel_threads(dispose_of, namespace):
     client = ndb.Client(namespace=namespace)
@@ -307,13 +303,12 @@ def insert(foo):
             entity = SomeKind(foo=foo, bar="none")
 
             key = entity.put()
+            dispose_of(key._key)
 
             retrieved = key.get()
             assert retrieved.foo == foo
             assert retrieved.bar == "none"
 
-            dispose_of(key._key)
-
     thread1 = threading.Thread(target=insert, args=[42], name="one")
     thread2 = threading.Thread(target=insert, args=[144], name="two")
 
@@ -332,12 +327,11 @@ class SomeKind(ndb.Model):
     foo = {str(i): i for i in range(500)}
     entity = SomeKind(foo=foo)
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == foo
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_compressed_json_property(dispose_of, ds_client):
@@ -347,12 +341,11 @@ class SomeKind(ndb.Model):
     foo = {str(i): i for i in range(500)}
     entity = SomeKind(foo=foo)
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == foo
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_compressed_blob_property(dispose_of, ds_client):
@@ -362,12 +355,11 @@ class SomeKind(ndb.Model):
     foo = b"abc" * 100
     entity = SomeKind(foo=foo)
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == foo
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_retrieve_entity_with_legacy_compressed_property(
@@ -399,12 +391,11 @@ class SomeKind(ndb.Model):
     foo = {str(i): i for i in range(500)}
     entity = SomeKind(foo=foo)
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == foo
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_key_property(dispose_of, ds_client):
@@ -414,14 +405,13 @@ class SomeKind(ndb.Model):
     key_value = ndb.Key("Whatevs", 123)
     entity = SomeKind(foo=key_value)
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == key_value
 
-    dispose_of(key._key)
 
-
-def test_insert_entity_with_caching(dispose_of, client_context):
+def test_insert_entity_with_caching(client_context):
     class SomeKind(ndb.Model):
         foo = ndb.IntegerProperty()
         bar = ndb.StringProperty()
@@ -454,6 +444,7 @@ class SomeKind(ndb.Model):
 
         entity = SomeKind(foo=42, bar="none")
         key = entity.put()
+        dispose_of(key._key)
         cache_key = _cache.global_cache_key(key._key)
         assert not cache_dict
 
@@ -470,8 +461,6 @@ class SomeKind(ndb.Model):
         # entity on write rather than waiting for a subsequent lookup.
         assert cache_key not in cache_dict
 
-        dispose_of(key._key)
-
 
 @pytest.mark.skipif(not USE_REDIS_CACHE, reason="Redis is not configured")
 def test_insert_entity_with_redis_cache(dispose_of, client_context):
@@ -485,6 +474,7 @@ class SomeKind(ndb.Model):
 
         entity = SomeKind(foo=42, bar="none")
         key = entity.put()
+        dispose_of(key._key)
         cache_key = _cache.global_cache_key(key._key)
         assert global_cache.redis.get(cache_key) is None
 
@@ -501,8 +491,6 @@ class SomeKind(ndb.Model):
         # entity on write rather than waiting for a subsequent lookup.
         assert global_cache.redis.get(cache_key) is None
 
-        dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_update_entity(ds_entity):
@@ -760,9 +748,8 @@ class SomeKind(ndb.Model):
     assert SomeKind.get_by_id(name) is None
 
     entity = SomeKind.get_or_insert(name, foo=21)
-    assert entity.foo == 21
-
     dispose_of(entity._key._key)
+    assert entity.foo == 21
 
 
 @pytest.mark.usefixtures("client_context")
@@ -793,6 +780,7 @@ class SomeKind(ndb.Model):
 
     entity = SomeKind(foo=42, bar=OtherKind(one="hi", two="mom"))
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == 42
@@ -801,8 +789,6 @@ class SomeKind(ndb.Model):
 
     assert isinstance(retrieved.bar, OtherKind)
 
-    dispose_of(key._key)
-
 
 def test_insert_entity_with_structured_property_legacy_data(
     client_context, dispose_of, ds_client
@@ -818,6 +804,7 @@ class SomeKind(ndb.Model):
     with client_context.new(legacy_data=True).use():
         entity = SomeKind(foo=42, bar=OtherKind(one="hi", two="mom"))
         key = entity.put()
+        dispose_of(key._key)
 
         retrieved = key.get()
         assert retrieved.foo == 42
@@ -831,8 +818,6 @@ class SomeKind(ndb.Model):
         assert ds_entity["bar.one"] == "hi"
         assert ds_entity["bar.two"] == "mom"
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_retrieve_entity_with_legacy_structured_property(ds_entity):
@@ -895,13 +880,12 @@ class SomeKind(ndb.Expando):
     entity = SomeKind(foo=42)
     entity.expando_prop = "exp-value"
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
     assert retrieved.foo == 42
     assert retrieved.expando_prop == "exp-value"
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_insert_polymodel(dispose_of):
@@ -916,6 +900,7 @@ class Cat(Feline):
 
     entity = Cat(one="hello", two="dad", three="i'm in jail")
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
 
@@ -925,8 +910,6 @@ class Cat(Feline):
     assert retrieved.two == "dad"
     assert retrieved.three == "i'm in jail"
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_insert_autonow_property(dispose_of):
@@ -937,14 +920,13 @@ class SomeKind(ndb.Model):
 
     entity = SomeKind(foo="bar")
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
 
     assert isinstance(retrieved.created_at, datetime.datetime)
     assert isinstance(retrieved.updated_at, datetime.datetime)
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_insert_nested_autonow_property(dispose_of):
@@ -957,14 +939,13 @@ class SomeKind(ndb.Model):
 
     entity = SomeKind(other=OtherKind())
     key = entity.put()
+    dispose_of(key._key)
 
     retrieved = key.get()
 
     assert isinstance(retrieved.other.created_at, datetime.datetime)
     assert isinstance(retrieved.other.updated_at, datetime.datetime)
 
-    dispose_of(key._key)
-
 
 @pytest.mark.usefixtures("client_context")
 def test_uninitialized_property(dispose_of):