fix(google): add support for Datastore emulator (#508)

Expands the google module with a DatastoreContainer using the beta
Datastore emulator using the same image as the PubSubContainer.

Im already using a local copy of this in production. It would be nice to
not have to support copy paste solutions and instead see it added to the
google module. This is my first PR so please let me know what I need to
do to get this over the line. Thanks.

Looks like @tillahoffmann wrote the original PubSub emulator container

---------

Co-authored-by: Matt Oates <[email protected]>
Co-authored-by: David Ankin <[email protected]>

Author: 3 people

modules/google/README.rst

@@ -1,2 +1,4 @@
+.. autoclass:: testcontainers.google.DatastoreContainer
+.. title:: testcontainers.google.DatastoreContainer
 .. autoclass:: testcontainers.google.PubSubContainer
 .. title:: testcontainers.google.PubSubContainer

modules/google/testcontainers/google/__init__.py

@@ -1 +1,2 @@
+from .datastore import DatastoreContainer  # noqa: F401
 from .pubsub import PubSubContainer  # noqa: F401

modules/google/testcontainers/google/datastore.py

@@ -0,0 +1,68 @@
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, 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 os
+from unittest.mock import patch
+
+from google.cloud import datastore
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_for_logs
+
+
+class DatastoreContainer(DockerContainer):
+    """
+    Datastore container for testing managed message queues.
+
+    Example:
+
+        The example will spin up a Google Cloud Datastore emulator that you can use for integration
+        tests. The :code:`datastore` instance provides convenience methods :code:`get_datastore_client` to
+        connect to the emulator without having to set the environment variable :code:`DATASTORE_EMULATOR_HOST`.
+
+        .. doctest::
+
+            >>> from testcontainers.google import DatastoreContainer
+
+            >>> config = DatastoreContainer()
+            >>> with config as datastore:
+            ...    datastore_client = datastore.get_datastore_client()
+    """
+
+    def __init__(
+        self,
+        image: str = "google/cloud-sdk:emulators",
+        project: str = "test-project",
+        port: int = 8081,
+        **kwargs,
+    ) -> None:
+        super().__init__(image=image, **kwargs)
+        self.project = project
+        self.port = port
+        self.with_exposed_ports(self.port)
+        self.with_command(
+            f"gcloud beta emulators datastore start --no-store-on-disk --project={project} --host-port=0.0.0.0:{port}"
+        )
+
+    def get_datastore_emulator_host(self) -> str:
+        return f"{self.get_container_host_ip()}:{self.get_exposed_port(self.port)}"
+
+    def get_datastore_client(self, **kwargs) -> datastore.Client:
+        wait_for_logs(self, "Dev App Server is now running.", timeout=30.0)
+        env_vars = {
+            "DATASTORE_DATASET": self.project,
+            "DATASTORE_EMULATOR_HOST": self.get_datastore_emulator_host(),
+            "DATASTORE_EMULATOR_HOST_PATH": f"{self.get_datastore_emulator_host()}/datastore",
+            "DATASTORE_HOST": f"http://{self.get_datastore_emulator_host()}",
+            "DATASTORE_PROJECT_ID": self.project,
+        }
+        with patch.dict(os.environ, env_vars):
+            return datastore.Client(**kwargs)

modules/google/tests/test_google.py

@@ -1,7 +1,8 @@
 from queue import Queue
+from google.cloud.datastore import Entity
 
 from testcontainers.core.waiting_utils import wait_for_logs
-from testcontainers.google import PubSubContainer
+from testcontainers.google import PubSubContainer, DatastoreContainer
 
 
 def test_pubsub_container():
@@ -27,3 +28,49 @@ def test_pubsub_container():
         message = queue.get(timeout=1)
         assert message.data == b"Hello world!"
         message.ack()
+
+
+def test_datastore_container_creation():
+    # Initialize the Datastore emulator container
+    with DatastoreContainer() as datastore:
+        # Obtain a datastore client configured to connect to the emulator
+        client = datastore.get_datastore_client()
+
+        # Define a unique key for a test entity to ensure test isolation
+        key = client.key("TestKind", "test_id_1")
+
+        # Create and insert a new entity
+        entity = Entity(key=key)
+        entity.update({"foo": "bar"})
+        client.put(entity)
+
+        # Fetch the just-inserted entity directly
+        fetched_entity = client.get(key)
+
+        # Assert that the fetched entity matches what was inserted
+        assert fetched_entity is not None, "Entity was not found in the datastore."
+        assert fetched_entity["foo"] == "bar", "Entity attribute 'foo' did not match expected value 'bar'."
+
+
+def test_datastore_container_isolation():
+    # Initialize the Datastore emulator container
+    with DatastoreContainer() as datastore:
+        # Obtain a datastore client configured to connect to the emulator
+        client = datastore.get_datastore_client()
+
+        # Define a unique key for a test entity to ensure test isolation
+        key = client.key("TestKind", "test_id_1")
+
+        # Create and insert a new entity
+        entity = Entity(key=key)
+        entity.update({"foo": "bar"})
+        client.put(entity)
+
+        # Create a second container and try to fetch the entity to makesure its a different container
+        with DatastoreContainer() as datastore2:
+            assert (
+                datastore.get_datastore_emulator_host() != datastore2.get_datastore_emulator_host()
+            ), "Datastore containers use the same port."
+            client2 = datastore2.get_datastore_client()
+            fetched_entity2 = client2.get(key)
+            assert fetched_entity2 is None, "Entity was found in the datastore."

poetry.lock

@@ -67,6 +67,7 @@ python-arango = { version = "^7.8", optional = true }
 azure-storage-blob = { version = "^12.19", optional = true }
 clickhouse-driver = { version = "*", optional = true }
 google-cloud-pubsub = { version = ">=2", optional = true }
+google-cloud-datastore = { version = ">=2", optional = true }
 influxdb = { version = "*", optional = true }
 influxdb-client = { version = "*", optional = true }
 kubernetes = { version = "*", optional = true }
@@ -90,7 +91,7 @@ arangodb = ["python-arango"]
 azurite = ["azure-storage-blob"]
 clickhouse = ["clickhouse-driver"]
 elasticsearch = []
-google = ["google-cloud-pubsub"]
+google = ["google-cloud-pubsub", "google-cloud-datastore"]
 influxdb = ["influxdb", "influxdb-client"]
 k3s = ["kubernetes", "pyyaml"]
 kafka = []