added tests

Author: daniel-sanche

tests/system/test_system.py

@@ -46,6 +46,7 @@
     TEST_DATABASES,
     TEST_DATABASES_W_ENTERPRISE,
     IS_KOKORO_TEST,
+    FIRESTORE_ENTERPRISE_DB,
 )
 
 
@@ -1688,6 +1689,54 @@ def test_query_stream_w_read_time(query_docs, cleanup, database):
     assert new_values[new_ref.id] == new_data
 
 
+@pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True)
+def test_pipeline_w_read_time(query_docs, cleanup, database):
+    # return early on kokoro. Test project doesn't currently support pipelines
+    # TODO: enable pipeline verification when kokoro test project is whitelisted
+    if IS_KOKORO_TEST:
+        pytest.skip("skipping pipeline verification on kokoro")
+
+    collection, stored, allowed_vals = query_docs
+    num_vals = len(allowed_vals)
+
+
+    # Find a read_time before adding the new document.
+    snapshots = collection.get()
+    read_time = snapshots[0].read_time
+
+    new_data = {
+        "a": 9000,
+        "b": 1,
+        "c": [10000, 1000],
+        "stats": {"sum": 9001, "product": 9000},
+    }
+    _, new_ref = collection.add(new_data)
+    # Add to clean-up.
+    cleanup(new_ref.delete)
+    stored[new_ref.id] = new_data
+
+    pipeline = collection.where(filter=FieldFilter("b", "==", 1)).pipeline()
+    # new query should have new_data
+    new_results = list(pipeline.stream())
+    new_values = {result.ref.id: result.data() for result in new_results}
+    assert len(new_values) == num_vals + 1
+    assert new_ref.id in new_values
+    assert new_values[new_ref.id] == new_data
+
+    # query with read_time should not have new)data
+    results = list(pipeline.stream(read_time=read_time))
+
+    values = {result.ref.id: result.data() for result in results}
+
+    assert len(values) == num_vals
+    assert new_ref.id not in values
+    for key, value in values.items():
+        assert stored[key] == value
+        assert value["b"] == 1
+        assert value["a"] != 9000
+        assert key != new_ref.id
+
+
 @pytest.mark.parametrize("database", TEST_DATABASES, indirect=True)
 def test_query_with_order_dot_key(client, cleanup, database):
     db = client

tests/system/test_system_async.py

@@ -57,6 +57,7 @@
     TEST_DATABASES,
     TEST_DATABASES_W_ENTERPRISE,
     IS_KOKORO_TEST,
+    FIRESTORE_ENTERPRISE_DB,
 )
 
 RETRIES = retries.AsyncRetry(
@@ -1611,6 +1612,54 @@ async def test_query_stream_w_read_time(query_docs, cleanup, database):
     assert new_values[new_ref.id] == new_data
 
 
+@pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True)
+async def test_pipeline_w_read_time(query_docs, cleanup, database):
+    # return early on kokoro. Test project doesn't currently support pipelines
+    # TODO: enable pipeline verification when kokoro test project is whitelisted
+    if IS_KOKORO_TEST:
+        pytest.skip("skipping pipeline verification on kokoro")
+
+    collection, stored, allowed_vals = query_docs
+    num_vals = len(allowed_vals)
+
+    # Find a read_time before adding the new document.
+    snapshots = await collection.get()
+    read_time = snapshots[0].read_time
+
+    new_data = {
+        "a": 9000,
+        "b": 1,
+        "c": [10000, 1000],
+        "stats": {"sum": 9001, "product": 9000},
+    }
+    _, new_ref = await collection.add(new_data)
+    # Add to clean-up.
+    cleanup(new_ref.delete)
+    stored[new_ref.id] = new_data
+
+    pipeline = collection.where(filter=FieldFilter("b", "==", 1)).pipeline()
+
+    # new query should have new_data
+    new_results = [result async for result in pipeline.stream()]
+    new_values = {result.ref.id: result.data() for result in new_results}
+    assert len(new_values) == num_vals + 1
+    assert new_ref.id in new_values
+    assert new_values[new_ref.id] == new_data
+
+    # pipeline with read_time should not have new_data
+    results = [result async for result in pipeline.stream(read_time=read_time)]
+
+    values = {result.ref.id: result.data() for result in results}
+
+    assert len(values) == num_vals
+    assert new_ref.id not in values
+    for key, value in values.items():
+        assert stored[key] == value
+        assert value["b"] == 1
+        assert value["a"] != 9000
+        assert key != new_ref.id
+
+
 @pytest.mark.parametrize("database", TEST_DATABASES, indirect=True)
 async def test_query_with_order_dot_key(client, cleanup, database):
     db = client

tests/unit/v1/test_async_pipeline.py

@@ -320,6 +320,36 @@ async def test_async_pipeline_stream_with_transaction():
     assert request.transaction == b"123"
 
 
+@pytest.mark.asyncio
+async def test_async_pipeline_stream_with_read_time():
+    """
+    test stream pipeline with read_time
+    """
+    import datetime
+
+    from google.cloud.firestore_v1.types import ExecutePipelineResponse
+    from google.cloud.firestore_v1.types import ExecutePipelineRequest
+
+    client = mock.Mock()
+    client.project = "A"
+    client._database = "B"
+    mock_rpc = mock.AsyncMock()
+    client._firestore_api.execute_pipeline = mock_rpc
+
+    read_time = datetime.datetime.now(tz=datetime.timezone.utc)
+
+    mock_rpc.return_value = _async_it([ExecutePipelineResponse()])
+    ppl_1 = _make_async_pipeline(client=client)
+
+    [r async for r in ppl_1.stream(read_time=read_time)]
+    assert mock_rpc.call_count == 1
+    request = mock_rpc.call_args[0][0]
+    assert isinstance(request, ExecutePipelineRequest)
+    assert request.structured_pipeline == ppl_1._to_pb()
+    assert request.database == "projects/A/databases/B"
+    assert request.read_time == read_time
+
+
 @pytest.mark.asyncio
 async def test_async_pipeline_stream_stream_equivalence():
     """
@@ -364,16 +394,22 @@ async def test_async_pipeline_stream_stream_equivalence_mocked():
     """
     pipeline.stream should call pipeline.stream internally
     """
+    import datetime
+
     ppl_1 = _make_async_pipeline()
     expected_data = [object(), object()]
-    expected_arg = object()
+    expected_transaction = object()
+    expected_read_time = datetime.datetime.now(tz=datetime.timezone.utc)
     with mock.patch.object(ppl_1, "stream") as mock_stream:
         mock_stream.return_value = _async_it(expected_data)
-        stream_results = await ppl_1.execute(expected_arg)
+        stream_results = await ppl_1.execute(
+            transaction=expected_transaction, read_time=expected_read_time
+        )
         assert mock_stream.call_count == 1
         assert mock_stream.call_args[0] == ()
-        assert len(mock_stream.call_args[1]) == 1
-        assert mock_stream.call_args[1]["transaction"] == expected_arg
+        assert len(mock_stream.call_args[1]) == 2
+        assert mock_stream.call_args[1]["transaction"] == expected_transaction
+        assert mock_stream.call_args[1]["read_time"] == expected_read_time
         assert stream_results == expected_data
 
 

tests/unit/v1/test_pipeline.py

@@ -302,6 +302,34 @@ def test_pipeline_stream_with_transaction():
     assert request.transaction == b"123"
 
 
+def test_pipeline_stream_with_read_time():
+    """
+    test stream pipeline with read_time
+    """
+    import datetime
+
+    from google.cloud.firestore_v1.types import ExecutePipelineResponse
+    from google.cloud.firestore_v1.types import ExecutePipelineRequest
+
+    client = mock.Mock()
+    client.project = "A"
+    client._database = "B"
+    mock_rpc = client._firestore_api.execute_pipeline
+
+    read_time = datetime.datetime.now(tz=datetime.timezone.utc)
+
+    mock_rpc.return_value = [ExecutePipelineResponse()]
+    ppl_1 = _make_pipeline(client=client)
+
+    list(ppl_1.stream(read_time=read_time))
+    assert mock_rpc.call_count == 1
+    request = mock_rpc.call_args[0][0]
+    assert isinstance(request, ExecutePipelineRequest)
+    assert request.structured_pipeline == ppl_1._to_pb()
+    assert request.database == "projects/A/databases/B"
+    assert request.read_time == read_time
+
+
 def test_pipeline_execute_stream_equivalence():
     """
     Pipeline.execute should provide same results from pipeline.stream, as a list
@@ -341,16 +369,22 @@ def test_pipeline_execute_stream_equivalence_mocked():
     """
     pipeline.execute should call pipeline.stream internally
     """
+    import datetime
+
     ppl_1 = _make_pipeline()
     expected_data = [object(), object()]
-    expected_arg = object()
+    expected_transaction = object()
+    expected_read_time = datetime.datetime.now(tz=datetime.timezone.utc)
     with mock.patch.object(ppl_1, "stream") as mock_stream:
         mock_stream.return_value = expected_data
-        stream_results = ppl_1.execute(expected_arg)
+        stream_results = ppl_1.execute(
+            transaction=expected_transaction, read_time=expected_read_time
+        )
         assert mock_stream.call_count == 1
         assert mock_stream.call_args[0] == ()
-        assert len(mock_stream.call_args[1]) == 1
-        assert mock_stream.call_args[1]["transaction"] == expected_arg
+        assert len(mock_stream.call_args[1]) == 2
+        assert mock_stream.call_args[1]["transaction"] == expected_transaction
+        assert mock_stream.call_args[1]["read_time"] == expected_read_time
         assert stream_results == expected_data