diff --git a/modules/google/testcontainers/google/__init__.py b/modules/google/testcontainers/google/__init__.py index b28f2ed48..c83d46234 100644 --- a/modules/google/testcontainers/google/__init__.py +++ b/modules/google/testcontainers/google/__init__.py @@ -1 +1,2 @@ from .pubsub import PubSubContainer # noqa: F401 +from .bigquery import BigQueryContainer # noqa: F401 diff --git a/modules/google/testcontainers/google/bigquery.py b/modules/google/testcontainers/google/bigquery.py new file mode 100644 index 000000000..5c9f34777 --- /dev/null +++ b/modules/google/testcontainers/google/bigquery.py @@ -0,0 +1,64 @@ +# +# 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. +from testcontainers.core.container import DockerContainer +from google.cloud.bigquery import Client as BigQueryClient +from google.api_core.client_options import ClientOptions +from google.auth.credentials import AnonymousCredentials + + +class BigQueryContainer(DockerContainer): + """ + Docker container to emulate BigQuery, based on https://github.com/testcontainers/testcontainers-java/blob/main/modules/gcloud/src/main/java/org/testcontainers/containers/BigQueryEmulatorContainer.java. + Uses ghcr.io/goccy/bigquery-emulator image by default. + Example: + + The example will spin up a Google Cloud BigQuery emulator that you can use for integration + tests. + + .. doctest:: + + >>> from testcontainers.google import BigQueryContainer + >>> from testcontainers.core.waiting_utils import wait_for_logs + >>> from google.cloud.bigquery import QueryJobConfig + + >>> with BigQueryContainer() as bigquery: + ... wait_for_logs(bigquery, "gRPC server listening", timeout=60) + ... client = bigquery.get_client() + ... result = client.query("SELECT 1", job_config=QueryJobConfig()).result() + ... print(result.total_rows) + """ + def __init__(self, image: str = "ghcr.io/goccy/bigquery-emulator:latest", project: str = "test-project", + http_port: int = 9050, grpc_port: int = 9060, **kwargs) -> None: + super(BigQueryContainer, self).__init__(image=image, **kwargs) + self.project = project + self.http_port = http_port + self.grpc_port = grpc_port + self.with_exposed_ports(http_port, grpc_port) + command = [ + "--project", project, + "--port", str(http_port), + "--grpc-port", str(grpc_port), + ] + self.with_command(' '.join(command)) + + def get_emulator_http_endpoint(self) -> str: + return f"http://{self.get_container_host_ip()}:{self.get_exposed_port(self.http_port)}" + + def get_client(self) -> BigQueryClient: + client_options = ClientOptions(api_endpoint=self.get_emulator_http_endpoint()) + return BigQueryClient( + project=self.project, + client_options=client_options, + credentials=AnonymousCredentials(), + ) + diff --git a/modules/google/tests/test_bigquery.py b/modules/google/tests/test_bigquery.py new file mode 100644 index 000000000..3e646da81 --- /dev/null +++ b/modules/google/tests/test_bigquery.py @@ -0,0 +1,36 @@ +from testcontainers.google import BigQueryContainer +from testcontainers.core.waiting_utils import wait_for_logs +from google.cloud.bigquery import QueryJobConfig, Client as BigQueryClient + + +def test_pubsub_container(): + with BigQueryContainer() as bigquery: + wait_for_logs(bigquery, "gRPC server listening", timeout=60) + + client: BigQueryClient = bigquery.get_client() + + # Function DDL + fn_stmt = ''' + CREATE FUNCTION testr(arr ARRAY>) AS ( + ( + SELECT SUM(IF(elem.name = "foo",elem.val,null)) + FROM UNNEST(arr) AS elem + ) + ) + ''' + + client.query(fn_stmt, job_config=QueryJobConfig()).result() + + select_stmt = ''' + SELECT + testr([ + STRUCT("foo", 10), + STRUCT("bar", 40), + STRUCT("foo", 20) + ]) + ''' + + result = client.query(select_stmt, job_config=QueryJobConfig()).result() + result = [ column for row in result for column in row ] + + assert result == [ 30 ] \ No newline at end of file diff --git a/modules/google/tests/test_pubsub.py b/modules/google/tests/test_pubsub.py new file mode 100644 index 000000000..780f5fdd6 --- /dev/null +++ b/modules/google/tests/test_pubsub.py @@ -0,0 +1,29 @@ +from queue import Queue + +from testcontainers.core.waiting_utils import wait_for_logs +from testcontainers.google import PubSubContainer + + +def test_pubsub_container(): + pubsub: PubSubContainer + with PubSubContainer() as pubsub: + wait_for_logs(pubsub, r"Server started, listening on \d+", timeout=60) + # Create a new topic + publisher = pubsub.get_publisher_client() + topic_path = publisher.topic_path(pubsub.project, "my-topic") + publisher.create_topic(name=topic_path) + + # Create a subscription + subscriber = pubsub.get_subscriber_client() + subscription_path = subscriber.subscription_path(pubsub.project, "my-subscription") + subscriber.create_subscription(name=subscription_path, topic=topic_path) + + # Publish a message + publisher.publish(topic_path, b"Hello world!") + + # Receive the message + queue = Queue() + subscriber.subscribe(subscription_path, queue.put) + message = queue.get(timeout=1) + assert message.data == b"Hello world!" + message.ack() diff --git a/pyproject.toml b/pyproject.toml index 7afb4cd96..09ed060ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ wrapt = "*" # "^1.16.0" python-arango = { version = "^7.8", optional = true } azure-storage-blob = { version = "^12.19", optional = true } clickhouse-driver = { version = "*", optional = true } +google-cloud-bigquery = { version = ">=2", optional = true } google-cloud-pubsub = { version = ">=2", optional = true } kubernetes = { version = "*", optional = true } pyyaml = { version = "*", optional = true } @@ -89,7 +90,7 @@ arangodb = ["python-arango"] azurite = ["azure-storage-blob"] clickhouse = ["clickhouse-driver"] elasticsearch = [] -google = ["google-cloud-pubsub"] +google = ["google-cloud-bigquery", "google-cloud-pubsub"] k3s = ["kubernetes", "pyyaml"] kafka = ["kafka-python"] keycloak = ["python-keycloak"]