wehub-resource-sync
95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterator
|
|
|
|
RERUN_DRAFT_PATH = str(Path(__file__).parent)
|
|
|
|
if RERUN_DRAFT_PATH not in sys.path:
|
|
sys.path.insert(0, RERUN_DRAFT_PATH)
|
|
|
|
import rerun_draft as rr # noqa: E402
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def populated_client(simple_dataset_prefix: Path) -> Iterator[rr.catalog.CatalogClient]:
|
|
"""Create a temporary dataset prefix with a few simple recordings."""
|
|
|
|
with rr.server.Server() as server:
|
|
client = server.client()
|
|
|
|
ds = client.create_dataset("basic_dataset")
|
|
ds.register_prefix(simple_dataset_prefix.as_uri())
|
|
|
|
# TODO(jleibs): Consider attaching this metadata table directly to the dataset
|
|
# and automatically joining it by default
|
|
meta = client.create_table(
|
|
"basic_dataset_metadata",
|
|
pa.schema([
|
|
("rerun_segment_id", pa.string()),
|
|
("success", pa.bool_()),
|
|
]),
|
|
)
|
|
|
|
meta.append(
|
|
rerun_segment_id=["simple_recording_0", "simple_recording_2"],
|
|
success=[True, False],
|
|
)
|
|
|
|
yield client
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def populated_client_complex(complex_dataset_prefix: Path) -> Iterator[rr.catalog.CatalogClient]:
|
|
"""Create a temporary dataset prefix with a few simple recordings."""
|
|
|
|
with rr.server.Server() as server:
|
|
client = server.client()
|
|
|
|
ds = client.create_dataset("complex_dataset")
|
|
ds.register_prefix(complex_dataset_prefix.as_uri())
|
|
|
|
# TODO(jleibs): Consider attaching this metadata table directly to the dataset
|
|
# and automatically joining it by default
|
|
meta = client.create_table(
|
|
"complex_dataset_metadata",
|
|
pa.schema([
|
|
("rerun_segment_id", pa.string()),
|
|
("success", pa.bool_()),
|
|
]),
|
|
)
|
|
|
|
meta.append(
|
|
rerun_segment_id=["complex_recording_1", "complex_recording_2", "complex_recording_3"],
|
|
success=[True, False, True],
|
|
)
|
|
|
|
yield client
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def basic_dataset(populated_client: rr.catalog.CatalogClient) -> Iterator[rr.catalog.DatasetEntry]:
|
|
yield populated_client.get_dataset(name="basic_dataset")
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def basic_metadata(populated_client: rr.catalog.CatalogClient) -> Iterator[rr.catalog.TableEntry]:
|
|
yield populated_client.get_table(name="basic_dataset_metadata")
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def complex_dataset(populated_client_complex: rr.catalog.CatalogClient) -> Iterator[rr.catalog.DatasetEntry]:
|
|
yield populated_client_complex.get_dataset(name="complex_dataset")
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def complex_metadata(populated_client_complex: rr.catalog.CatalogClient) -> Iterator[rr.catalog.TableEntry]:
|
|
yield populated_client_complex.get_table(name="complex_dataset_metadata")
|