Refactor exercise functions to fixtures

tests/datastore_firestore/test_async_batching.py

@@ -12,22 +12,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
+
 from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
 from newrelic.api.background_task import background_task
 from testing_support.validators.validate_database_duration import (
     validate_database_duration,
 )
 
-async def _exercise_async_write_batch(async_client, async_collection):
-    docs = [async_collection.document(str(x)) for x in range(1, 4)]
-    async_batch = async_client.batch()
-    for doc in docs:
-        async_batch.set(doc, {})
 
-    await async_batch.commit()
+@pytest.fixture()
+def exercise_async_write_batch(async_client, async_collection):
+    async def _exercise_async_write_batch():
+        docs = [async_collection.document(str(x)) for x in range(1, 4)]
+        async_batch = async_client.batch()
+        for doc in docs:
+            async_batch.set(doc, {})
+
+        await async_batch.commit()
+    return _exercise_async_write_batch
 
 
-def test_firestore_async_write_batch(loop, async_client, async_collection):
+def test_firestore_async_write_batch(loop, exercise_async_write_batch):
     _test_scoped_metrics = [
         ("Datastore/operation/Firestore/commit", 1),
     ]
@@ -45,6 +51,6 @@ def test_firestore_async_write_batch(loop, async_client, async_collection):
     )
     @background_task(name="test_firestore_async_write_batch")
     def _test():
-        loop.run_until_complete(_exercise_async_write_batch(async_client, async_collection))
+        loop.run_until_complete(exercise_async_write_batch())
 
     _test()

tests/datastore_firestore/test_async_client.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import pytest
 
 from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
@@ -27,13 +28,16 @@ def existing_document(collection):
     return doc
 
 
-async def _exercise_async_client(async_client, existing_document):
-    assert len([_ async for _ in async_client.collections()]) >= 1
-    doc = [_ async for _ in async_client.get_all([existing_document])][0]
-    assert doc.to_dict()["x"] == 1
+@pytest.fixture()
+def exercise_async_client(async_client, existing_document):
+    async def _exercise_async_client():
+        assert len([_ async for _ in async_client.collections()]) >= 1
+        doc = [_ async for _ in async_client.get_all([existing_document])][0]
+        assert doc.to_dict()["x"] == 1
+    return _exercise_async_client
 
 
-def test_firestore_async_client(loop, async_client, existing_document):
+def test_firestore_async_client(loop, exercise_async_client):
     _test_scoped_metrics = [
         ("Datastore/operation/Firestore/collections", 1),
         ("Datastore/operation/Firestore/get_all", 1),
@@ -53,7 +57,7 @@ def test_firestore_async_client(loop, async_client, existing_document):
     )
     @background_task(name="test_firestore_async_client")
     def _test():
-       loop.run_until_complete(_exercise_async_client(async_client, existing_document))
+       loop.run_until_complete(exercise_async_client())
 
     _test()
 

tests/datastore_firestore/test_async_collections.py

@@ -12,27 +12,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
+
 from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
 from newrelic.api.background_task import background_task
 from testing_support.validators.validate_database_duration import (
     validate_database_duration,
 )
 
 
-async def _exercise_collections(async_collection):
-    async_collection.document("DoesNotExist")
-    await async_collection.add({"capital": "Rome", "currency": "Euro", "language": "Italian"}, "Italy")
-    await async_collection.add({"capital": "Mexico City", "currency": "Peso", "language": "Spanish"}, "Mexico")
-
-    documents_get = await async_collection.get()
-    assert len(documents_get) == 2
-    documents_stream = [_ async for _ in async_collection.stream()]
-    assert len(documents_stream) == 2
-    documents_list = [_ async for _ in async_collection.list_documents()]
-    assert len(documents_list) == 2
+@pytest.fixture()
+def exercise_collections(async_collection):
+    async def _exercise_collections():
+        async_collection.document("DoesNotExist")
+        await async_collection.add({"capital": "Rome", "currency": "Euro", "language": "Italian"}, "Italy")
+        await async_collection.add({"capital": "Mexico City", "currency": "Peso", "language": "Spanish"}, "Mexico")
+
+        documents_get = await async_collection.get()
+        assert len(documents_get) == 2
+        documents_stream = [_ async for _ in async_collection.stream()]
+        assert len(documents_stream) == 2
+        documents_list = [_ async for _ in async_collection.list_documents()]
+        assert len(documents_list) == 2
+    return _exercise_collections
 
 
-def test_firestore_async_collections(loop, async_collection):
+def test_firestore_async_collections(loop, exercise_collections, async_collection):
     _test_scoped_metrics = [
         ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1),
         ("Datastore/statement/Firestore/%s/get" % async_collection.id, 1),
@@ -57,13 +62,13 @@ def test_firestore_async_collections(loop, async_collection):
     )
     @background_task(name="test_firestore_async_collections")
     def _test():
-        loop.run_until_complete(_exercise_collections(async_collection))
+        loop.run_until_complete(exercise_collections())
 
     _test()
 
 
 @background_task()
-def test_firestore_async_collections_generators(loop, collection, async_collection, assert_trace_for_async_generator):
+def test_firestore_async_collections_generators(collection, async_collection, assert_trace_for_async_generator):
     collection.add({})
     collection.add({})
     assert len([_ for _ in collection.list_documents()]) == 2

tests/datastore_firestore/test_async_documents.py

@@ -12,30 +12,35 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
+
 from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
 from newrelic.api.background_task import background_task
 from testing_support.validators.validate_database_duration import (
     validate_database_duration,
 )
 
 
-async def _exercise_async_documents(async_collection):
-    italy_doc = async_collection.document("Italy")
-    await italy_doc.set({"capital": "Rome", "currency": "Euro", "language": "Italian"})
-    await italy_doc.get()
-    italian_cities = italy_doc.collection("cities")
-    await italian_cities.add({"capital": "Rome"})
-    retrieved_coll = [_ async for _ in italy_doc.collections()]
-    assert len(retrieved_coll) == 1
+@pytest.fixture()
+def exercise_async_documents(async_collection):
+    async def _exercise_async_documents():
+        italy_doc = async_collection.document("Italy")
+        await italy_doc.set({"capital": "Rome", "currency": "Euro", "language": "Italian"})
+        await italy_doc.get()
+        italian_cities = italy_doc.collection("cities")
+        await italian_cities.add({"capital": "Rome"})
+        retrieved_coll = [_ async for _ in italy_doc.collections()]
+        assert len(retrieved_coll) == 1
 
-    usa_doc = async_collection.document("USA")
-    await usa_doc.create({"capital": "Washington D.C.", "currency": "Dollar", "language": "English"})
-    await usa_doc.update({"president": "Joe Biden"})
+        usa_doc = async_collection.document("USA")
+        await usa_doc.create({"capital": "Washington D.C.", "currency": "Dollar", "language": "English"})
+        await usa_doc.update({"president": "Joe Biden"})
 
-    await async_collection.document("USA").delete()
+        await async_collection.document("USA").delete()
+    return _exercise_async_documents
 
 
-def test_firestore_async_documents(loop, async_collection):
+def test_firestore_async_documents(loop, exercise_async_documents):
     _test_scoped_metrics = [
         ("Datastore/statement/Firestore/Italy/set", 1),
         ("Datastore/statement/Firestore/Italy/get", 1),
@@ -66,7 +71,7 @@ def test_firestore_async_documents(loop, async_collection):
     )
     @background_task(name="test_firestore_async_documents")
     def _test():
-        loop.run_until_complete(_exercise_async_documents(async_collection))
+        loop.run_until_complete(exercise_async_documents())
 
     _test()
 

tests/datastore_firestore/test_async_query.py

@@ -33,13 +33,16 @@ def sample_data(collection):
 
 # ===== AsyncQuery =====
 
-async def _exercise_async_query(async_collection):
-    async_query = async_collection.select("x").limit(10).order_by("x").where(field_path="x", op_string="<=", value=3)
-    assert len(await async_query.get()) == 3
-    assert len([_ async for _ in async_query.stream()]) == 3
+@pytest.fixture()
+def exercise_async_query(async_collection):
+    async def _exercise_async_query():
+        async_query = async_collection.select("x").limit(10).order_by("x").where(field_path="x", op_string="<=", value=3)
+        assert len(await async_query.get()) == 3
+        assert len([_ async for _ in async_query.stream()]) == 3
+    return _exercise_async_query
 
 
-def test_firestore_async_query(loop, async_collection):
+def test_firestore_async_query(loop, exercise_async_query, async_collection):
     _test_scoped_metrics = [
         ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1),
         ("Datastore/statement/Firestore/%s/get" % async_collection.id, 1),
@@ -60,7 +63,7 @@ def test_firestore_async_query(loop, async_collection):
     )
     @background_task(name="test_firestore_async_query")
     def _test():
-        loop.run_until_complete(_exercise_async_query(async_collection))
+        loop.run_until_complete(exercise_async_query())
 
     _test()
 
@@ -72,13 +75,16 @@ def test_firestore_async_query_generators(async_collection, assert_trace_for_asy
 
 # ===== AsyncAggregationQuery =====
 
-async def _exercise_async_aggregation_query(async_collection):
-    async_aggregation_query = async_collection.select("x").where(field_path="x", op_string="<=", value=3).count()
-    assert (await async_aggregation_query.get())[0][0].value == 3
-    assert [_ async for _ in async_aggregation_query.stream()][0][0].value == 3
+@pytest.fixture()
+def exercise_async_aggregation_query(async_collection):
+    async def _exercise_async_aggregation_query():
+        async_aggregation_query = async_collection.select("x").where(field_path="x", op_string="<=", value=3).count()
+        assert (await async_aggregation_query.get())[0][0].value == 3
+        assert [_ async for _ in async_aggregation_query.stream()][0][0].value == 3
+    return _exercise_async_aggregation_query
 
 
-def test_firestore_async_aggregation_query(loop, async_collection):
+def test_firestore_async_aggregation_query(loop, exercise_async_aggregation_query, async_collection):
     _test_scoped_metrics = [
         ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1),
         ("Datastore/statement/Firestore/%s/get" % async_collection.id, 1),
@@ -99,7 +105,7 @@ def test_firestore_async_aggregation_query(loop, async_collection):
     )
     @background_task(name="test_firestore_async_aggregation_query")
     def _test():
-        loop.run_until_complete(_exercise_async_aggregation_query(async_collection))
+        loop.run_until_complete(exercise_async_aggregation_query())
 
     _test()
 
@@ -138,20 +144,23 @@ async def _mock_partition_query():
     yield
 
 
-async def _exercise_async_collection_group(async_client, async_collection):
-    async_collection_group = async_client.collection_group(async_collection.id)
-    assert len(await async_collection_group.get())
-    assert len([d async for d in async_collection_group.stream()])
-
-    partitions = [p async for p in async_collection_group.get_partitions(1)]
-    assert len(partitions) == 2
-    documents = []
-    while partitions:
-        documents.extend(await partitions.pop().query().get())
-    assert len(documents) == 6
-
-
-def test_firestore_async_collection_group(loop, async_client, async_collection, patch_partition_queries):
+@pytest.fixture()
+def exercise_async_collection_group(async_client, async_collection):
+    async def _exercise_async_collection_group():
+        async_collection_group = async_client.collection_group(async_collection.id)
+        assert len(await async_collection_group.get())
+        assert len([d async for d in async_collection_group.stream()])
+
+        partitions = [p async for p in async_collection_group.get_partitions(1)]
+        assert len(partitions) == 2
+        documents = []
+        while partitions:
+            documents.extend(await partitions.pop().query().get())
+        assert len(documents) == 6
+    return _exercise_async_collection_group
+
+
+def test_firestore_async_collection_group(loop, exercise_async_collection_group, async_collection, patch_partition_queries):
     _test_scoped_metrics = [
         ("Datastore/statement/Firestore/%s/get" % async_collection.id, 3),
         ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1),
@@ -175,7 +184,7 @@ def test_firestore_async_collection_group(loop, async_client, async_collection,
     )
     @background_task(name="test_firestore_async_collection_group")
     def _test():
-        loop.run_until_complete(_exercise_async_collection_group(async_client, async_collection))
+        loop.run_until_complete(exercise_async_collection_group())
 
     _test()