added tests

Author: daniel-sanche

tests/unit/v1/test_async_pipeline.py

@@ -15,6 +15,8 @@
 import mock
 import pytest
 
+from google.cloud.firestore_v1 import pipeline_stages as stages
+
 
 def _make_async_pipeline(*args, client=mock.Mock()):
     from google.cloud.firestore_v1.async_pipeline import AsyncPipeline
@@ -54,11 +56,9 @@ def test_async_pipeline_repr_single_stage():
 
 
 def test_async_pipeline_repr_multiple_stage():
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage, Collection
-
-    stage_1 = Collection("path")
-    stage_2 = GenericStage("second", 2)
-    stage_3 = GenericStage("third", 3)
+    stage_1 = stages.Collection("path")
+    stage_2 = stages.GenericStage("second", 2)
+    stage_3 = stages.GenericStage("third", 3)
     ppl = _make_async_pipeline(stage_1, stage_2, stage_3)
     repr_str = repr(ppl)
     assert repr_str == (
@@ -71,10 +71,8 @@ def test_async_pipeline_repr_multiple_stage():
 
 
 def test_async_pipeline_repr_long():
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage
-
     num_stages = 100
-    stage_list = [GenericStage("custom", i) for i in range(num_stages)]
+    stage_list = [stages.GenericStage("custom", i) for i in range(num_stages)]
     ppl = _make_async_pipeline(*stage_list)
     repr_str = repr(ppl)
     assert repr_str.count("GenericStage") == num_stages
@@ -83,10 +81,9 @@ def test_async_pipeline_repr_long():
 
 def test_async_pipeline__to_pb():
     from google.cloud.firestore_v1.types.pipeline import StructuredPipeline
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage
 
-    stage_1 = GenericStage("first")
-    stage_2 = GenericStage("second")
+    stage_1 = stages.GenericStage("first")
+    stage_2 = stages.GenericStage("second")
     ppl = _make_async_pipeline(stage_1, stage_2)
     pb = ppl._to_pb()
     assert isinstance(pb, StructuredPipeline)
@@ -96,11 +93,9 @@ def test_async_pipeline__to_pb():
 
 def test_async_pipeline_append():
     """append should create a new pipeline with the additional stage"""
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage
-
-    stage_1 = GenericStage("first")
+    stage_1 = stages.GenericStage("first")
     ppl_1 = _make_async_pipeline(stage_1, client=object())
-    stage_2 = GenericStage("second")
+    stage_2 = stages.GenericStage("second")
     ppl_2 = ppl_1._append(stage_2)
     assert ppl_1 != ppl_2
     assert len(ppl_1.stages) == 1
@@ -118,15 +113,14 @@ async def test_async_pipeline_execute_empty():
     """
     from google.cloud.firestore_v1.types import ExecutePipelineResponse
     from google.cloud.firestore_v1.types import ExecutePipelineRequest
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage
 
     client = mock.Mock()
     client.project = "A"
     client._database = "B"
     mock_rpc = mock.AsyncMock()
     client._firestore_api.execute_pipeline = mock_rpc
     mock_rpc.return_value = _async_it([ExecutePipelineResponse()])
-    ppl_1 = _make_async_pipeline(GenericStage("s"), client=client)
+    ppl_1 = _make_async_pipeline(stages.GenericStage("s"), client=client)
 
     results = [r async for r in ppl_1.execute()]
     assert results == []
@@ -145,7 +139,6 @@ async def test_async_pipeline_execute_no_doc_ref():
     from google.cloud.firestore_v1.types import Document
     from google.cloud.firestore_v1.types import ExecutePipelineResponse
     from google.cloud.firestore_v1.types import ExecutePipelineRequest
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage
     from google.cloud.firestore_v1.pipeline_result import PipelineResult
 
     client = mock.Mock()
@@ -156,7 +149,7 @@ async def test_async_pipeline_execute_no_doc_ref():
     mock_rpc.return_value = _async_it(
         [ExecutePipelineResponse(results=[Document()], execution_time={"seconds": 9})]
     )
-    ppl_1 = _make_async_pipeline(GenericStage("s"), client=client)
+    ppl_1 = _make_async_pipeline(stages.GenericStage("s"), client=client)
 
     results = [r async for r in ppl_1.execute()]
     assert len(results) == 1
@@ -315,3 +308,20 @@ async def test_async_pipeline_execute_with_transaction():
     assert request.structured_pipeline == ppl_1._to_pb()
     assert request.database == "projects/A/databases/B"
     assert request.transaction == b"123"
+
+@pytest.mark.parametrize("method,args,result_cls", [
+    ("select", (), stages.Select),
+    ("where", (mock.Mock(),), stages.Where),
+    ("sort", (), stages.Sort),
+    ("offset", (1,), stages.Offset),
+    ("limit", (1,), stages.Limit),
+
+])
+def test_async_pipeline_methods(method, args, result_cls):
+    start_ppl = _make_async_pipeline()
+    method_ptr = getattr(start_ppl, method)
+    result_ppl = method_ptr(*args)
+    assert result_ppl != start_ppl
+    assert len(start_ppl.stages) == 0
+    assert len(result_ppl.stages) == 1
+    assert isinstance(result_ppl.stages[0], result_cls)

tests/unit/v1/test_pipeline.py

@@ -13,6 +13,9 @@
 # limitations under the License
 
 import mock
+import pytest
+
+from google.cloud.firestore_v1 import pipeline_stages as stages
 
 
 def _make_pipeline(*args, client=mock.Mock()):
@@ -48,11 +51,9 @@ def test_pipeline_repr_single_stage():
 
 
 def test_pipeline_repr_multiple_stage():
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage, Collection
-
-    stage_1 = Collection("path")
-    stage_2 = GenericStage("second", 2)
-    stage_3 = GenericStage("third", 3)
+    stage_1 = stages.Collection("path")
+    stage_2 = stages.GenericStage("second", 2)
+    stage_3 = stages.GenericStage("third", 3)
     ppl = _make_pipeline(stage_1, stage_2, stage_3)
     repr_str = repr(ppl)
     assert repr_str == (
@@ -65,10 +66,8 @@ def test_pipeline_repr_multiple_stage():
 
 
 def test_pipeline_repr_long():
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage
-
     num_stages = 100
-    stage_list = [GenericStage("custom", i) for i in range(num_stages)]
+    stage_list = [stages.GenericStage("custom", i) for i in range(num_stages)]
     ppl = _make_pipeline(*stage_list)
     repr_str = repr(ppl)
     assert repr_str.count("GenericStage") == num_stages
@@ -77,10 +76,9 @@ def test_pipeline_repr_long():
 
 def test_pipeline__to_pb():
     from google.cloud.firestore_v1.types.pipeline import StructuredPipeline
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage
 
-    stage_1 = GenericStage("first")
-    stage_2 = GenericStage("second")
+    stage_1 = stages.GenericStage("first")
+    stage_2 = stages.GenericStage("second")
     ppl = _make_pipeline(stage_1, stage_2)
     pb = ppl._to_pb()
     assert isinstance(pb, StructuredPipeline)
@@ -90,11 +88,9 @@ def test_pipeline__to_pb():
 
 def test_pipeline_append():
     """append should create a new pipeline with the additional stage"""
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage
-
-    stage_1 = GenericStage("first")
+    stage_1 = stages.GenericStage("first")
     ppl_1 = _make_pipeline(stage_1, client=object())
-    stage_2 = GenericStage("second")
+    stage_2 = stages.GenericStage("second")
     ppl_2 = ppl_1._append(stage_2)
     assert ppl_1 != ppl_2
     assert len(ppl_1.stages) == 1
@@ -111,14 +107,13 @@ def test_pipeline_execute_empty():
     """
     from google.cloud.firestore_v1.types import ExecutePipelineResponse
     from google.cloud.firestore_v1.types import ExecutePipelineRequest
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage
 
     client = mock.Mock()
     client.project = "A"
     client._database = "B"
     mock_rpc = client._firestore_api.execute_pipeline
     mock_rpc.return_value = [ExecutePipelineResponse()]
-    ppl_1 = _make_pipeline(GenericStage("s"), client=client)
+    ppl_1 = _make_pipeline(stages.GenericStage("s"), client=client)
 
     results = list(ppl_1.execute())
     assert results == []
@@ -136,7 +131,6 @@ def test_pipeline_execute_no_doc_ref():
     from google.cloud.firestore_v1.types import Document
     from google.cloud.firestore_v1.types import ExecutePipelineResponse
     from google.cloud.firestore_v1.types import ExecutePipelineRequest
-    from google.cloud.firestore_v1.pipeline_stages import GenericStage
     from google.cloud.firestore_v1.pipeline_result import PipelineResult
 
     client = mock.Mock()
@@ -146,7 +140,7 @@ def test_pipeline_execute_no_doc_ref():
     mock_rpc.return_value = [
         ExecutePipelineResponse(results=[Document()], execution_time={"seconds": 9})
     ]
-    ppl_1 = _make_pipeline(GenericStage("s"), client=client)
+    ppl_1 = _make_pipeline(stages.GenericStage("s"), client=client)
 
     results = list(ppl_1.execute())
     assert len(results) == 1
@@ -295,3 +289,20 @@ def test_pipeline_execute_with_transaction():
     assert request.structured_pipeline == ppl_1._to_pb()
     assert request.database == "projects/A/databases/B"
     assert request.transaction == b"123"
+
+@pytest.mark.parametrize("method,args,result_cls", [
+    ("select", (), stages.Select),
+    ("where", (mock.Mock(),), stages.Where),
+    ("sort", (), stages.Sort),
+    ("offset", (1,), stages.Offset),
+    ("limit", (1,), stages.Limit),
+
+])
+def test_pipeline_methods(method, args, result_cls):
+    start_ppl = _make_pipeline()
+    method_ptr = getattr(start_ppl, method)
+    result_ppl = method_ptr(*args)
+    assert result_ppl != start_ppl
+    assert len(start_ppl.stages) == 0
+    assert len(result_ppl.stages) == 1
+    assert isinstance(result_ppl.stages[0], result_cls)

tests/unit/v1/test_pipeline_source.py

@@ -44,6 +44,14 @@ def test_collection(self):
         assert isinstance(first_stage, stages.Collection)
         assert first_stage.path == "/path"
 
+    def test_collection_group(self):
+        instance = self._make_client().pipeline()
+        ppl = instance.collection_group("id")
+        assert isinstance(ppl, self._expected_pipeline_type)
+        assert len(ppl.stages) == 1
+        first_stage = ppl.stages[0]
+        assert isinstance(first_stage, stages.CollectionGroup)
+        assert first_stage.collection_id == "id"
 
 class TestPipelineSourceWithAsyncClient(TestPipelineSource):
     """