fix: Improve Firestore tests by separating test runs, improving cleanup (#2667)

* fix: harden the wait for consistency for firestore multiple listen snippet

* fix: formatting

* add a setup step in CLI (#2611)

* fix: formatting

* fix: formatting

* fix: add setup to snippets for tests

* fix: rename sydney_query to denver_query as it refers to denver

* fix: Use stream instead of get in samples

* fix: stop deleting LA, delete DEN instead

* fix: move delete tests to end of file

* fix: move BJ to top to work with snapshot cursors

* fix: Beijing -> BJ

* during delete collection delete all collections created for these tests

* fix: add cleanup test to distributed counters

* wrap db.collections calls to use unique names

* fix: remove unused import

* fix: remove more complicated fix for time now that we have separate collections

* fix: reorder tests

* fix: move test firestore client to test file

* fix: lint

* fix: use old-style classes for Python 2.7

Co-authored-by: Tianzi Cai <[email protected]>

Author: 2 people

firestore/cloud-client/distributed_counters_test.py

@@ -47,3 +47,15 @@ def test_distributed_counters(fs_client):
     counter.increment_counter(doc_ref)
     counter.increment_counter(doc_ref)
     assert counter.get_count(doc_ref) == 2
+
+
+def test_distributed_counters_cleanup(fs_client):
+    col = fs_client.collection("dc_samples")
+    doc_ref = col.document("distributed_counter")
+
+    shards = doc_ref.collection("shards").list_documents()
+    shards_list = [shard for shard in shards]
+    for shard in shards_list:
+        shard.delete()
+
+    doc_ref.delete()

firestore/cloud-client/snippets.py

@@ -159,6 +159,8 @@ def add_example_data():
     db = firestore.Client()
     # [START add_example_data]
     cities_ref = db.collection(u'cities')
+    cities_ref.document(u'BJ').set(
+        City(u'Beijing', None, u'China', True, 21500000, [u'hebei']).to_dict())
     cities_ref.document(u'SF').set(
         City(u'San Francisco', u'CA', u'USA', False, 860000,
              [u'west_coast', u'norcal']).to_dict())
@@ -171,8 +173,6 @@ def add_example_data():
     cities_ref.document(u'TOK').set(
         City(u'Tokyo', None, u'Japan', True, 9000000,
              [u'kanto', u'honshu']).to_dict())
-    cities_ref.document(u'BJ').set(
-        City(u'Beijing', None, u'China', True, 21500000, [u'hebei']).to_dict())
     # [END add_example_data]
 
 
@@ -304,6 +304,10 @@ def structure_subcollection_ref():
 
 def update_doc():
     db = firestore.Client()
+    db.collection(u'cities').document(u'DC').set(
+        City(u'Washington D.C.', None, u'USA', True, 680000,
+             [u'east_coast']).to_dict())
+
     # [START update_doc]
     city_ref = db.collection(u'cities').document(u'DC')
 
@@ -314,6 +318,10 @@ def update_doc():
 
 def update_doc_array():
     db = firestore.Client()
+    db.collection(u'cities').document(u'DC').set(
+        City(u'Washington D.C.', None, u'USA', True, 680000,
+             [u'east_coast']).to_dict())
+
     # [START fs_update_doc_array]
     city_ref = db.collection(u'cities').document(u'DC')
 
@@ -329,6 +337,10 @@ def update_doc_array():
 
 def update_multiple():
     db = firestore.Client()
+    db.collection(u'cities').document(u'DC').set(
+        City(u'Washington D.C.', None, u'USA', True, 680000,
+             [u'east_coast']).to_dict())
+
     # [START update_multiple]
     doc_ref = db.collection(u'cities').document(u'DC')
 
@@ -441,9 +453,9 @@ def update_data_batch():
     sf_ref = db.collection(u'cities').document(u'SF')
     batch.update(sf_ref, {u'population': 1000000})
 
-    # Delete LA
-    la_ref = db.collection(u'cities').document(u'LA')
-    batch.delete(la_ref)
+    # Delete DEN
+    den_ref = db.collection(u'cities').document(u'DEN')
+    batch.delete(den_ref)
 
     # Commit the batch
     batch.commit()
@@ -490,12 +502,12 @@ def compound_query_valid_multi_clause():
     # [START compound_query_valid_multi_clause]
     cities_ref = db.collection(u'cities')
 
-    sydney_query = cities_ref.where(
+    denver_query = cities_ref.where(
         u'state', u'==', u'CO').where(u'name', u'==', u'Denver')
     large_us_cities_query = cities_ref.where(
         u'state', u'==', u'CA').where(u'population', u'>', 1000000)
     # [END compound_query_valid_multi_clause]
-    print(sydney_query)
+    print(denver_query)
     print(large_us_cities_query)
 
 
@@ -703,6 +715,7 @@ def on_snapshot(col_snapshot, changes, read_time):
     }
     db.collection(u'cities').document(u'SF').set(data)
     sleep(1)
+
     query_watch.unsubscribe()
 
 
@@ -737,6 +750,7 @@ def on_snapshot(col_snapshot, changes, read_time):
         u'capital': False,
         u'population': 80000
     })
+    sleep(1)
 
     # Modifying document
     mtv_document.update({
@@ -746,6 +760,7 @@ def on_snapshot(col_snapshot, changes, read_time):
         u'capital': False,
         u'population': 90000
     })
+    sleep(1)
 
     # Delete document
     mtv_document.delete()
@@ -801,7 +816,7 @@ def delete_full_collection():
 
     # [START delete_full_collection]
     def delete_collection(coll_ref, batch_size):
-        docs = coll_ref.limit(batch_size).get()
+        docs = coll_ref.limit(batch_size).stream()
         deleted = 0
 
         for doc in docs:
@@ -814,6 +829,9 @@ def delete_collection(coll_ref, batch_size):
     # [END delete_full_collection]
 
     delete_collection(db.collection(u'cities'), 10)
+    delete_collection(db.collection(u'data'), 10)
+    delete_collection(db.collection(u'objects'), 10)
+    delete_collection(db.collection(u'users'), 10)
 
 
 def collection_group_query(db):

firestore/cloud-client/snippets_test.py

@@ -12,6 +12,7 @@
 # limitations under the License.
 
 import os
+import uuid
 
 from google.cloud import firestore
 import pytest
@@ -20,10 +21,26 @@
 
 os.environ['GOOGLE_CLOUD_PROJECT'] = os.environ['FIRESTORE_PROJECT']
 
+UNIQUE_STRING = str(uuid.uuid4()).split("-")[0]
+
+
+class TestFirestoreClient(firestore.Client):
+    def __init__(self, *args, **kwargs):
+        self._UNIQUE_STRING = UNIQUE_STRING
+        self._super = super(TestFirestoreClient, self)
+        self._super.__init__(*args, **kwargs)
+
+    def collection(self, collection_name, *args, **kwargs):
+        collection_name += '-{}'.format(self._UNIQUE_STRING)
+        return self._super.collection(collection_name, *args, **kwargs)
+
+
+snippets.firestore.Client = TestFirestoreClient
+
 
 @pytest.fixture
 def db():
-    yield firestore.Client()
+    yield snippets.firestore.Client()
 
 
 def test_quickstart_new_instance():
@@ -250,15 +267,6 @@ def test_cursor_multiple_conditions():
     snippets.cursor_multiple_conditions()
 
 
-def test_delete_single_doc():
-    snippets.delete_single_doc()
-
-
-def test_delete_field(db):
-    db.collection('cities').document('Beijing').set({'capital': True})
-    snippets.delete_field()
-
-
 def test_listen_document(capsys):
     snippets.listen_document()
     out, _ = capsys.readouterr()
@@ -280,6 +288,15 @@ def test_listen_for_changes(capsys):
     assert 'Removed city: MTV' in out
 
 
+def test_delete_single_doc():
+    snippets.delete_single_doc()
+
+
+def test_delete_field(db):
+    db.collection('cities').document('BJ').set({'capital': True})
+    snippets.delete_field()
+
+
 def test_delete_full_collection():
     snippets.delete_full_collection()