Skip to content

Commit 5983d84

Browse files
committed
1 parent 7517297 commit 5983d84

File tree

7 files changed

+2467
-1774
lines changed

7 files changed

+2467
-1774
lines changed

conf.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# testcontainers documentation build configuration file, created by
1+
# testcontainers documentation build configuration file, created by
22
# sphinx-quickstart on Tue Mar 21 21:09:48 2017.
33
#
44
# This file is execfile()d with the current directory set to its
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
from .datastore import DatastoreContainer # noqa: F401
22
from .pubsub import PubSubContainer # noqa: F401
3+
from .bigquery import BigQueryContainer # noqa: F401
Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
#
2+
# Licensed under the Apache License, Version 2.0 (the "License"); you may
3+
# not use this file except in compliance with the License. You may obtain
4+
# a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
10+
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
11+
# License for the specific language governing permissions and limitations
12+
# under the License.
13+
from testcontainers.core.container import DockerContainer
14+
from google.cloud.bigquery import Client as BigQueryClient
15+
from google.api_core.client_options import ClientOptions
16+
from google.auth.credentials import AnonymousCredentials
17+
18+
19+
class BigQueryContainer(DockerContainer):
20+
"""
21+
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.
22+
Uses ghcr.io/goccy/bigquery-emulator image by default.
23+
Example:
24+
25+
The example will spin up a Google Cloud BigQuery emulator that you can use for integration
26+
tests.
27+
28+
.. doctest::
29+
30+
>>> from testcontainers.google import BigQueryContainer
31+
>>> from testcontainers.core.waiting_utils import wait_for_logs
32+
>>> from google.cloud.bigquery import QueryJobConfig
33+
34+
>>> with BigQueryContainer(platform="linux/x86_64") as bigquery:
35+
... wait_for_logs(bigquery, "gRPC server listening", timeout=60)
36+
... client = bigquery.get_client()
37+
... result = client.query("SELECT 1", job_config=QueryJobConfig()).result()
38+
... print(result.total_rows)
39+
"""
40+
def __init__(self, image: str = "ghcr.io/goccy/bigquery-emulator:latest", project: str = "test-project",
41+
http_port: int = 9050, grpc_port: int = 9060, **kwargs) -> None:
42+
super(BigQueryContainer, self).__init__(image=image, **kwargs)
43+
self.project = project
44+
self.http_port = http_port
45+
self.grpc_port = grpc_port
46+
self.with_exposed_ports(http_port, grpc_port)
47+
command = [
48+
"--project", project,
49+
"--port", str(http_port),
50+
"--grpc-port", str(grpc_port),
51+
]
52+
self.with_command(' '.join(command))
53+
54+
def get_emulator_http_endpoint(self) -> str:
55+
return f"http://{self.get_container_host_ip()}:{self.get_exposed_port(self.http_port)}"
56+
57+
def get_client(self) -> BigQueryClient:
58+
client_options = ClientOptions(api_endpoint=self.get_emulator_http_endpoint())
59+
return BigQueryClient(
60+
project=self.project,
61+
client_options=client_options,
62+
credentials=AnonymousCredentials(),
63+
)
64+

modules/google/tests/test_bigquery.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import os
2+
3+
from testcontainers.google import BigQueryContainer
4+
from testcontainers.core.waiting_utils import wait_for_logs
5+
from google.cloud.bigquery import QueryJobConfig, Client as BigQueryClient
6+
7+
8+
9+
def test_pubsub_container():
10+
with BigQueryContainer as bigquery:
11+
wait_for_logs(bigquery, "gRPC server listening", timeout=60)
12+
13+
client: BigQueryClient = bigquery.get_client()
14+
15+
# Function DDL
16+
fn_stmt = '''
17+
CREATE FUNCTION testr(arr ARRAY<STRUCT<name STRING, val INT64>>) AS (
18+
(
19+
SELECT SUM(IF(elem.name = "foo",elem.val,null))
20+
FROM UNNEST(arr) AS elem
21+
)
22+
)
23+
'''
24+
25+
client.query(fn_stmt, job_config=QueryJobConfig()).result()
26+
27+
select_stmt = '''
28+
SELECT
29+
testr([
30+
STRUCT<name STRING, val INT64>("foo", 10),
31+
STRUCT<name STRING, val INT64>("bar", 40),
32+
STRUCT<name STRING, val INT64>("foo", 20)
33+
])
34+
'''
35+
36+
result = client.query(select_stmt, job_config=QueryJobConfig()).result()
37+
result = [ column for row in result for column in row ]
38+
39+
assert result == [ 30 ]

modules/google/tests/test_pubsub.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from queue import Queue
2+
3+
from testcontainers.core.waiting_utils import wait_for_logs
4+
from testcontainers.google import PubSubContainer
5+
6+
7+
def test_pubsub_container():
8+
pubsub: PubSubContainer
9+
with PubSubContainer() as pubsub:
10+
wait_for_logs(pubsub, r"Server started, listening on \d+", timeout=60)
11+
# Create a new topic
12+
publisher = pubsub.get_publisher_client()
13+
topic_path = publisher.topic_path(pubsub.project, "my-topic")
14+
publisher.create_topic(name=topic_path)
15+
16+
# Create a subscription
17+
subscriber = pubsub.get_subscriber_client()
18+
subscription_path = subscriber.subscription_path(pubsub.project, "my-subscription")
19+
subscriber.create_subscription(name=subscription_path, topic=topic_path)
20+
21+
# Publish a message
22+
publisher.publish(topic_path, b"Hello world!")
23+
24+
# Receive the message
25+
queue = Queue()
26+
subscriber.subscribe(subscription_path, queue.put)
27+
message = queue.get(timeout=1)
28+
assert message.data == b"Hello world!"
29+
message.ack()

0 commit comments

Comments
 (0)