Updating datastore samples to use eventually consistent helper (#331)

Change-Id: I3a0404b354ff01530de07a6227a0f34b9aa9edcf

datastore/api/snippets_test.py

@@ -11,9 +11,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import time
-
 from gcloud import datastore
+from gcp.testing import eventually_consistent
 from gcp.testing.flaky import flaky
 import pytest
 import snippets
@@ -28,17 +27,8 @@ def __init__(self, *args, **kwargs):
     def cleanup(self):
         with self.batch():
             self.delete_multi(
-                [x.key for x in self.entities_to_delete] +
-                self.keys_to_delete)
-
-
-# This is pretty hacky, but make datastore wait 1s after any
-# put operation to in order to account for eventual consistency.
-class WaitingClient(CleanupClient):
-    def put_multi(self, *args, **kwargs):
-        result = super(WaitingClient, self).put_multi(*args, **kwargs)
-        time.sleep(1)
-        return result
+                list(set([x.key for x in self.entities_to_delete])) +
+                list(set(self.keys_to_delete)))
 
 
 @pytest.yield_fixture
@@ -48,13 +38,6 @@ def client(cloud_config):
     client.cleanup()
 
 
-@pytest.yield_fixture
-def waiting_client(cloud_config):
-    client = WaitingClient(cloud_config.project)
-    yield client
-    client.cleanup()
-
-
 @flaky
 class TestDatastoreSnippets:
     # These tests mostly just test the absence of exceptions.
@@ -118,20 +101,23 @@ def test_batch_lookup(self, client):
     def test_batch_delete(self, client):
         snippets.batch_delete(client)
 
-    def test_unindexed_property_query(self, waiting_client):
-        tasks = snippets.unindexed_property_query(waiting_client)
-        waiting_client.entities_to_delete.extend(tasks)
+    @eventually_consistent.mark
+    def test_unindexed_property_query(self, client):
+        tasks = snippets.unindexed_property_query(client)
+        client.entities_to_delete.extend(tasks)
         assert tasks
 
-    def test_basic_query(self, waiting_client):
-        tasks = snippets.basic_query(waiting_client)
-        waiting_client.entities_to_delete.extend(tasks)
+    @eventually_consistent.mark
+    def test_basic_query(self, client):
+        tasks = snippets.basic_query(client)
+        client.entities_to_delete.extend(tasks)
         assert tasks
 
-    def test_projection_query(self, waiting_client):
-        priorities, percents = snippets.projection_query(waiting_client)
-        waiting_client.entities_to_delete.extend(
-            waiting_client.query(kind='Task').fetch())
+    @eventually_consistent.mark
+    def test_projection_query(self, client):
+        priorities, percents = snippets.projection_query(client)
+        client.entities_to_delete.extend(
+            client.query(kind='Task').fetch())
         assert priorities
         assert percents
 
@@ -143,63 +129,74 @@ def test_ancestor_query(self, client):
     def test_run_query(self, client):
         snippets.run_query(client)
 
-    def test_cursor_paging(self, waiting_client):
+    def test_cursor_paging(self, client):
         for n in range(6):
-            waiting_client.entities_to_delete.append(
-                snippets.insert(waiting_client))
-
-        page_one, cursor_one, page_two, cursor_two = snippets.cursor_paging(
-            waiting_client)
-
-        assert len(page_one) == 5
-        assert len(page_two) == 1
-        assert cursor_one
-        assert cursor_two
-
-    def test_property_filter(self, waiting_client):
-        tasks = snippets.property_filter(waiting_client)
-        waiting_client.entities_to_delete.extend(tasks)
+            client.entities_to_delete.append(
+                snippets.insert(client))
+
+        @eventually_consistent.call
+        def _():
+            results = snippets.cursor_paging(client)
+            page_one, cursor_one, page_two, cursor_two = results
+
+            assert len(page_one) == 5
+            assert len(page_two) == 1
+            assert cursor_one
+            assert cursor_two
+
+    @eventually_consistent.mark
+    def test_property_filter(self, client):
+        tasks = snippets.property_filter(client)
+        client.entities_to_delete.extend(tasks)
         assert tasks
 
-    def test_composite_filter(self, waiting_client):
-        tasks = snippets.composite_filter(waiting_client)
-        waiting_client.entities_to_delete.extend(tasks)
+    @eventually_consistent.mark
+    def test_composite_filter(self, client):
+        tasks = snippets.composite_filter(client)
+        client.entities_to_delete.extend(tasks)
         assert tasks
 
-    def test_key_filter(self, waiting_client):
-        tasks = snippets.key_filter(waiting_client)
-        waiting_client.entities_to_delete.extend(tasks)
+    @eventually_consistent.mark
+    def test_key_filter(self, client):
+        tasks = snippets.key_filter(client)
+        client.entities_to_delete.extend(tasks)
         assert tasks
 
-    def test_ascending_sort(self, waiting_client):
-        tasks = snippets.ascending_sort(waiting_client)
-        waiting_client.entities_to_delete.extend(tasks)
+    @eventually_consistent.mark
+    def test_ascending_sort(self, client):
+        tasks = snippets.ascending_sort(client)
+        client.entities_to_delete.extend(tasks)
         assert tasks
 
-    def test_descending_sort(self, waiting_client):
-        tasks = snippets.descending_sort(waiting_client)
-        waiting_client.entities_to_delete.extend(tasks)
+    @eventually_consistent.mark
+    def test_descending_sort(self, client):
+        tasks = snippets.descending_sort(client)
+        client.entities_to_delete.extend(tasks)
         assert tasks
 
-    def test_multi_sort(self, waiting_client):
-        tasks = snippets.multi_sort(waiting_client)
-        waiting_client.entities_to_delete.extend(tasks)
+    @eventually_consistent.mark
+    def test_multi_sort(self, client):
+        tasks = snippets.multi_sort(client)
+        client.entities_to_delete.extend(tasks)
         assert tasks
 
-    def test_keys_only_query(self, waiting_client):
-        keys = snippets.keys_only_query(waiting_client)
-        waiting_client.entities_to_delete.extend(
-            waiting_client.query(kind='Task').fetch())
+    @eventually_consistent.mark
+    def test_keys_only_query(self, client):
+        keys = snippets.keys_only_query(client)
+        client.entities_to_delete.extend(
+            client.query(kind='Task').fetch())
         assert keys
 
-    def test_distinct_query(self, waiting_client):
-        tasks = snippets.distinct_query(waiting_client)
-        waiting_client.entities_to_delete.extend(tasks)
+    @eventually_consistent.mark
+    def test_distinct_query(self, client):
+        tasks = snippets.distinct_query(client)
+        client.entities_to_delete.extend(tasks)
         assert tasks
 
-    def test_distinct_on_query(self, waiting_client):
-        tasks = snippets.distinct_on_query(waiting_client)
-        waiting_client.entities_to_delete.extend(tasks)
+    @eventually_consistent.mark
+    def test_distinct_on_query(self, client):
+        tasks = snippets.distinct_on_query(client)
+        client.entities_to_delete.extend(tasks)
         assert tasks
 
     def test_kindless_query(self, client):
@@ -251,29 +248,33 @@ def transactional_single_entity_group_read_only(self, client):
         assert task_list
         assert tasks_in_list
 
-    def test_namespace_run_query(self, waiting_client):
+    @eventually_consistent.mark
+    def test_namespace_run_query(self, client):
         all_namespaces, filtered_namespaces = snippets.namespace_run_query(
-            waiting_client)
+            client)
         assert all_namespaces
         assert filtered_namespaces
         assert 'google' in filtered_namespaces
 
-    def test_kind_run_query(self, waiting_client):
-        kinds = snippets.kind_run_query(waiting_client)
-        waiting_client.entities_to_delete.extend(
-            waiting_client.query(kind='Task').fetch())
+    @eventually_consistent.mark
+    def test_kind_run_query(self, client):
+        kinds = snippets.kind_run_query(client)
+        client.entities_to_delete.extend(
+            client.query(kind='Task').fetch())
         assert kinds
         assert 'Task' in kinds
 
-    def test_property_run_query(self, waiting_client):
-        kinds = snippets.property_run_query(waiting_client)
-        waiting_client.entities_to_delete.extend(
-            waiting_client.query(kind='Task').fetch())
+    @eventually_consistent.mark
+    def test_property_run_query(self, client):
+        kinds = snippets.property_run_query(client)
+        client.entities_to_delete.extend(
+            client.query(kind='Task').fetch())
         assert kinds
         assert 'Task' in kinds
 
-    def test_property_by_kind_run_query(self, waiting_client):
-        reprs = snippets.property_by_kind_run_query(waiting_client)
-        waiting_client.entities_to_delete.extend(
-            waiting_client.query(kind='Task').fetch())
+    @eventually_consistent.mark
+    def test_property_by_kind_run_query(self, client):
+        reprs = snippets.property_by_kind_run_query(client)
+        client.entities_to_delete.extend(
+            client.query(kind='Task').fetch())
         assert reprs

datastore/api/tasks_test.py

@@ -12,6 +12,7 @@
 # limitations under the License.
 
 from gcloud import datastore
+from gcp.testing import eventually_consistent
 from gcp.testing.flaky import flaky
 import pytest
 import tasks
@@ -55,8 +56,11 @@ def test_mark_done(client):
 def test_list_tasks(client):
     task1_key = tasks.add_task(client, 'Test task 1')
     task2_key = tasks.add_task(client, 'Test task 2')
-    task_list = tasks.list_tasks(client)
-    assert [x.key for x in task_list] == [task1_key, task2_key]
+
+    @eventually_consistent.call
+    def _():
+        task_list = tasks.list_tasks(client)
+        assert [x.key for x in task_list] == [task1_key, task2_key]
 
 
 @flaky
@@ -72,9 +76,11 @@ def test_format_tasks(client):
     tasks.add_task(client, 'Test task 2')
     tasks.mark_done(client, task1_key.id)
 
-    output = tasks.format_tasks(tasks.list_tasks(client))
+    @eventually_consistent.call
+    def _():
+        output = tasks.format_tasks(tasks.list_tasks(client))
 
-    assert 'Test task 1' in output
-    assert 'Test task 2' in output
-    assert 'done' in output
-    assert 'created' in output
+        assert 'Test task 1' in output
+        assert 'Test task 2' in output
+        assert 'done' in output
+        assert 'created' in output