testcontainers · alexanderankin · Mar 30, 2024 · Mar 26, 2024 · Mar 30, 2024 · Mar 30, 2024
diff --git a/modules/google/README.rst b/modules/google/README.rst
@@ -1,2 +1,4 @@
+.. autoclass:: testcontainers.google.DatastoreContainer
+.. title:: testcontainers.google.DatastoreContainer
 .. autoclass:: testcontainers.google.PubSubContainer
 .. title:: testcontainers.google.PubSubContainer
diff --git a/modules/google/testcontainers/google/__init__.py b/modules/google/testcontainers/google/__init__.py
@@ -1 +1,2 @@
+from .datastore import DatastoreContainer  # noqa: F401
 from .pubsub import PubSubContainer  # noqa: F401
diff --git a/modules/google/testcontainers/google/datastore.py b/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)
diff --git a/modules/google/tests/test_google.py b/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."
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -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 = []
Original file line number	Diff line number	Diff line change
		@@ -1 +1,2 @@
		from .datastore import DatastoreContainer # noqa: F401
		from .pubsub import PubSubContainer # noqa: F401