Skip to content

Commit 800efe8

Browse files
provinzkrautprovinzkraut
committed
migrate bigquery
1 parent 12a74ba commit 800efe8

File tree

3 files changed

+76
-175
lines changed

3 files changed

+76
-175
lines changed

‎src/pytest_databases/docker/bigquery.py‎

Lines changed: 64 additions & 120 deletions
Original file line numberDiff line numberDiff line change
@@ -1,151 +1,95 @@
11
from __future__ import annotations
22

3-
import os
4-
import sys
5-
from pathlib import Path
3+
from dataclasses import dataclass
64
from typing import TYPE_CHECKING
75

86
import pytest
97
from google.api_core.client_options import ClientOptions
108
from google.auth.credentials import AnonymousCredentials, Credentials
119
from google.cloud import bigquery
1210

13-
from pytest_databases.docker import DockerServiceRegistry
14-
from pytest_databases.helpers import simple_string_hash
11+
from pytest_databases._service import DockerService
12+
from pytest_databases.helpers import get_xdist_worker_id
13+
from pytest_databases.types import ServiceContainer
1514

1615
if TYPE_CHECKING:
1716
from collections.abc import Generator
1817

1918

20-
COMPOSE_PROJECT_NAME: str = f"pytest-databases-bigquery-{simple_string_hash(__file__)}"
19+
@dataclass
20+
class BigQueryService(ServiceContainer):
21+
project: str
22+
dataset: str
23+
credentials: Credentials
2124

25+
@property
26+
def endpoint(self) -> str:
27+
return f"http://{self.host}:{self.port}"
2228

23-
def bigquery_responsive(
24-
host: str,
25-
bigquery_endpoint: str,
26-
bigquery_dataset: str,
27-
bigquery_client_options: ClientOptions,
28-
bigquery_project: str,
29-
bigquery_credentials: Credentials,
30-
) -> bool:
31-
try:
32-
client = bigquery.Client(
33-
project=bigquery_project, client_options=bigquery_client_options, credentials=bigquery_credentials
34-
)
35-
36-
job = client.query(query="SELECT 1 as one")
37-
38-
resp = list(job.result())
39-
return resp[0].one == 1
40-
except Exception: # noqa: BLE001
41-
return False
29+
@property
30+
def client_options(self) -> ClientOptions:
31+
return ClientOptions(api_endpoint=self.endpoint)
4232

4333

4434
@pytest.fixture(scope="session")
45-
def bigquery_compose_project_name() -> str:
46-
return os.environ.get("COMPOSE_PROJECT_NAME", COMPOSE_PROJECT_NAME)
47-
48-
49-
@pytest.fixture(autouse=False, scope="session")
50-
def bigquery_docker_services(
51-
bigquery_compose_project_name: str, worker_id: str = "main"
52-
) -> Generator[DockerServiceRegistry, None, None]:
53-
if os.getenv("GITHUB_ACTIONS") == "true" and sys.platform != "linux":
54-
pytest.skip("Docker not available on this platform")
55-
56-
with DockerServiceRegistry(worker_id, compose_project_name=bigquery_compose_project_name) as registry:
57-
yield registry
58-
59-
60-
@pytest.fixture(scope="session")
61-
def bigquery_port() -> int:
62-
return 9051
63-
64-
65-
@pytest.fixture(scope="session")
66-
def bigquery_grpc_port() -> int:
67-
return 9061
68-
69-
70-
@pytest.fixture(scope="session")
71-
def bigquery_dataset() -> str:
72-
return "test-dataset"
73-
74-
75-
@pytest.fixture(scope="session")
76-
def bigquery_project() -> str:
77-
return "emulator-test-project"
78-
79-
80-
@pytest.fixture(scope="session")
81-
def bigquery_client_options(bigquery_endpoint: str) -> ClientOptions:
82-
return ClientOptions(api_endpoint=bigquery_endpoint)
83-
84-
85-
@pytest.fixture(scope="session")
86-
def bigquery_credentials() -> Credentials:
87-
return AnonymousCredentials()
88-
89-
90-
@pytest.fixture(scope="session")
91-
def bigquery_docker_compose_files() -> list[Path]:
92-
return [Path(Path(__file__).parent / "docker-compose.bigquery.yml")]
93-
94-
95-
@pytest.fixture(scope="session")
96-
def default_bigquery_service_name() -> str:
97-
return "bigquery"
98-
99-
100-
@pytest.fixture(scope="session")
101-
def bigquery_docker_ip(bigquery_docker_services: DockerServiceRegistry) -> str:
102-
return bigquery_docker_services.docker_ip
103-
104-
105-
@pytest.fixture(scope="session")
106-
def bigquery_endpoint(bigquery_docker_ip: str, bigquery_port: int) -> str:
107-
return f"http://{bigquery_docker_ip}:{bigquery_port}"
35+
def bigquery_xdist_isolate() -> bool:
36+
return True
10837

10938

11039
@pytest.fixture(autouse=False, scope="session")
11140
def bigquery_service(
112-
bigquery_docker_services: DockerServiceRegistry,
113-
default_bigquery_service_name: str,
114-
bigquery_docker_compose_files: list[Path],
115-
bigquery_port: int,
116-
bigquery_grpc_port: int,
117-
bigquery_endpoint: str,
118-
bigquery_dataset: str,
119-
bigquery_project: str,
120-
bigquery_credentials: Credentials,
121-
bigquery_client_options: ClientOptions,
122-
) -> Generator[None, None, None]:
123-
os.environ["BIGQUERY_ENDPOINT"] = bigquery_endpoint
124-
os.environ["BIGQUERY_DATASET"] = bigquery_dataset
125-
os.environ["BIGQUERY_PORT"] = str(bigquery_port)
126-
os.environ["BIGQUERY_GRPC_PORT"] = str(bigquery_grpc_port)
127-
os.environ["GOOGLE_CLOUD_PROJECT"] = bigquery_project
128-
bigquery_docker_services.start(
129-
name=default_bigquery_service_name,
130-
docker_compose_files=bigquery_docker_compose_files,
41+
docker_service: DockerService,
42+
bigquery_xdist_isolate: bool,
43+
) -> Generator[BigQueryService, None, None]:
44+
project = "emulator-test-project"
45+
dataset = "test-dataset"
46+
47+
def check(_service: ServiceContainer) -> bool:
48+
try:
49+
client = bigquery.Client(
50+
project=project,
51+
client_options=ClientOptions(api_endpoint=f"http://{_service.host}:{_service.port}"),
52+
credentials=AnonymousCredentials(),
53+
)
54+
55+
job = client.query(query="SELECT 1 as one")
56+
57+
resp = list(job.result())
58+
return resp[0].one == 1
59+
except Exception: # noqa: BLE001
60+
return False
61+
62+
container_name = "bigquery"
63+
if not bigquery_xdist_isolate:
64+
container_name = f"{container_name}_{get_xdist_worker_id()}"
65+
66+
with docker_service.run(
67+
image="ghcr.io/goccy/bigquery-emulator:latest",
68+
command=f"--project={project} --dataset={dataset}",
69+
name=container_name,
70+
check=check,
71+
env={
72+
"PROJECT_ID": project,
73+
"DATASET_NAME": dataset,
74+
},
75+
container_port=9050,
13176
timeout=60,
132-
check=bigquery_responsive,
133-
bigquery_endpoint=bigquery_endpoint,
134-
bigquery_dataset=bigquery_dataset,
135-
bigquery_project=bigquery_project,
136-
bigquery_credentials=bigquery_credentials,
137-
bigquery_client_options=bigquery_client_options,
138-
)
139-
yield
77+
) as service:
78+
yield BigQueryService(
79+
host=service.host,
80+
port=service.port,
81+
project=project,
82+
dataset=dataset,
83+
credentials=AnonymousCredentials(),
84+
)
14085

14186

14287
@pytest.fixture(autouse=False, scope="session")
14388
def bigquery_startup_connection(
144-
bigquery_service: DockerServiceRegistry,
145-
bigquery_project: str,
146-
bigquery_credentials: Credentials,
147-
bigquery_client_options: ClientOptions,
89+
bigquery_service: BigQueryService,
14890
) -> Generator[bigquery.Client, None, None]:
14991
yield bigquery.Client(
150-
project=bigquery_project, client_options=bigquery_client_options, credentials=bigquery_credentials
92+
project=bigquery_service.project,
93+
client_options=bigquery_service.client_options,
94+
credentials=bigquery_service.credentials,
15195
)

‎src/pytest_databases/docker/docker-compose.bigquery.yml‎

Lines changed: 0 additions & 18 deletions
This file was deleted.

‎tests/docker/test_bigquery.py‎

Lines changed: 12 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -4,52 +4,27 @@
44

55
from google.cloud import bigquery
66

7-
from pytest_databases.docker.bigquery import bigquery_responsive
7+
from pytest_databases.docker.bigquery import BigQueryService
88

99
if TYPE_CHECKING:
10-
from google.api_core.client_options import ClientOptions
11-
from google.auth.credentials import Credentials
12-
13-
from pytest_databases.docker import DockerServiceRegistry
10+
pass
1411

1512
pytest_plugins = [
1613
"pytest_databases.docker.bigquery",
1714
]
1815

1916

20-
def test_bigquery_default_config(
21-
bigquery_docker_ip: str,
22-
bigquery_port: int,
23-
bigquery_grpc_port: int,
24-
bigquery_dataset: str,
25-
bigquery_endpoint: str,
26-
bigquery_project: str,
27-
) -> None:
28-
assert bigquery_port == 9051
29-
assert bigquery_grpc_port == 9061
30-
assert bigquery_dataset == "test-dataset"
31-
assert bigquery_endpoint == f"http://{bigquery_docker_ip}:9051"
32-
assert bigquery_project == "emulator-test-project"
33-
34-
35-
def test_bigquery_services(
36-
bigquery_docker_ip: str,
37-
bigquery_service: DockerServiceRegistry,
38-
bigquery_endpoint: str,
39-
bigquery_dataset: str,
40-
bigquery_client_options: ClientOptions,
41-
bigquery_project: str,
42-
bigquery_credentials: Credentials,
43-
) -> None:
44-
ping = bigquery_responsive(
45-
bigquery_docker_ip,
46-
bigquery_endpoint=bigquery_endpoint,
47-
bigquery_dataset=bigquery_dataset,
48-
bigquery_project=bigquery_project,
49-
bigquery_credentials=bigquery_credentials,
50-
bigquery_client_options=bigquery_client_options,
17+
def test_bigquery_service(bigquery_service: BigQueryService) -> None:
18+
client = bigquery.Client(
19+
project=bigquery_service.project,
20+
client_options=bigquery_service.client_options,
21+
credentials=bigquery_service.credentials,
5122
)
52-
assert ping
23+
24+
job = client.query(query="SELECT 1 as one")
25+
26+
resp = list(job.result())
27+
assert resp[0].one == 1
5328

5429

5530
def test_bigquery_service_after_start(

0 commit comments

Comments
 (0)