wehub-resource-sync
704 lines
24 KiB
Python
704 lines
24 KiB
Python
"""Tests for TurbopufferDatasink.
|
|
|
|
Organized by critical paths:
|
|
1. Constructor validation
|
|
2. Client initialization
|
|
3. Arrow table preparation
|
|
4. Single-namespace batching
|
|
5. Transform to Turbopuffer format
|
|
6. Retry logic
|
|
7. End-to-end write orchestration
|
|
8. Streaming behavior
|
|
9. Multi-namespace writes
|
|
10. Serialization
|
|
"""
|
|
|
|
import pickle
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from typing import List
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pyarrow as pa
|
|
import pytest
|
|
from packaging.version import parse as parse_version
|
|
|
|
from ray.data._internal.datasource.turbopuffer_datasink import TurbopufferDatasink
|
|
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
|
|
|
# Skip all tests if PyArrow version is less than 19.0
|
|
pytestmark = pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("19.0.0"),
|
|
reason="TurbopufferDatasink tests require PyArrow >= 19.0",
|
|
)
|
|
|
|
# =============================================================================
|
|
# Fixtures
|
|
# =============================================================================
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_turbopuffer_module(monkeypatch):
|
|
"""Provide a fake turbopuffer module so imports in the datasink succeed."""
|
|
fake_module = MagicMock()
|
|
fake_module.Turbopuffer = MagicMock()
|
|
with patch.dict(sys.modules, {"turbopuffer": fake_module}):
|
|
yield fake_module
|
|
|
|
|
|
@pytest.fixture
|
|
def sink():
|
|
"""Default sink with minimal required arguments."""
|
|
return TurbopufferDatasink(
|
|
namespace="default_ns",
|
|
region="gcp-us-central1",
|
|
api_key="test-api-key",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_client():
|
|
"""Mock Turbopuffer client with namespace support."""
|
|
client = MagicMock()
|
|
client.namespace.return_value = MagicMock()
|
|
return client
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_table():
|
|
"""Standard table with id and vector columns."""
|
|
return pa.table(
|
|
{
|
|
"id": [1, 2, 3],
|
|
"vector": [[0.1], [0.2], [0.3]],
|
|
}
|
|
)
|
|
|
|
|
|
def make_sink(**kwargs) -> TurbopufferDatasink:
|
|
"""Helper to construct a sink with minimal required arguments."""
|
|
params = {
|
|
"namespace": "default_ns",
|
|
"region": "gcp-us-central1",
|
|
"api_key": "test-api-key",
|
|
}
|
|
params.update(kwargs)
|
|
return TurbopufferDatasink(**params)
|
|
|
|
|
|
# =============================================================================
|
|
# 1. Constructor validation
|
|
# =============================================================================
|
|
|
|
|
|
class TestConstructorValidation:
|
|
"""Tests for constructor argument validation."""
|
|
|
|
def test_requires_namespace_or_namespace_column(self):
|
|
"""Must provide exactly one of namespace / namespace_column."""
|
|
with pytest.raises(ValueError, match="Either.*must be provided"):
|
|
TurbopufferDatasink(
|
|
region="gcp-us-central1",
|
|
api_key="k",
|
|
)
|
|
|
|
def test_rejects_both_namespace_and_namespace_column(self):
|
|
"""Cannot provide both namespace and namespace_column."""
|
|
with pytest.raises(ValueError, match="exactly one"):
|
|
TurbopufferDatasink(
|
|
namespace="ns",
|
|
namespace_column="ns_col",
|
|
region="gcp-us-central1",
|
|
api_key="k",
|
|
)
|
|
|
|
def test_namespace_column_cannot_be_id_or_vector(self):
|
|
"""namespace_column must not collide with id_column or vector_column."""
|
|
with pytest.raises(ValueError, match="namespace_column.*must not be the same"):
|
|
make_sink(namespace=None, namespace_column="id")
|
|
|
|
with pytest.raises(ValueError, match="namespace_column.*must not be the same"):
|
|
make_sink(namespace=None, namespace_column="vector")
|
|
|
|
def test_api_key_from_env(self, monkeypatch):
|
|
"""API key can come from environment variable."""
|
|
monkeypatch.delenv("TURBOPUFFER_API_KEY", raising=False)
|
|
|
|
# No api_key and no env var -> error
|
|
with pytest.raises(ValueError):
|
|
TurbopufferDatasink(namespace="ns", region="gcp-us-central1")
|
|
|
|
# With env var, init should succeed
|
|
monkeypatch.setenv("TURBOPUFFER_API_KEY", "env-api-key")
|
|
sink = TurbopufferDatasink(namespace="ns", region="gcp-us-central1")
|
|
assert sink.api_key == "env-api-key"
|
|
|
|
def test_rejects_same_id_and_vector_column(self):
|
|
"""id_column and vector_column must be distinct."""
|
|
with pytest.raises(ValueError, match="id_column and vector_column"):
|
|
make_sink(id_column="doc_id", vector_column="doc_id")
|
|
|
|
def test_accepts_region_only(self):
|
|
"""Constructor succeeds with region and no base_url."""
|
|
sink = make_sink(region="gcp-us-central1")
|
|
assert sink.region == "gcp-us-central1"
|
|
assert sink.base_url is None
|
|
|
|
def test_accepts_base_url_only(self):
|
|
"""Constructor succeeds with base_url and no region."""
|
|
sink = make_sink(
|
|
region=None,
|
|
base_url="https://gcp-us-central1.turbopuffer.com",
|
|
)
|
|
assert sink.base_url == "https://gcp-us-central1.turbopuffer.com"
|
|
assert sink.region is None
|
|
|
|
def test_rejects_both_region_and_base_url(self):
|
|
"""Cannot provide both region and base_url."""
|
|
with pytest.raises(ValueError, match="exactly one of 'region' or 'base_url'"):
|
|
make_sink(
|
|
region="gcp-us-central1",
|
|
base_url="https://gcp-us-central1.turbopuffer.com",
|
|
)
|
|
|
|
def test_rejects_neither_region_nor_base_url(self):
|
|
"""Must provide at least one of region or base_url."""
|
|
with pytest.raises(ValueError, match="Either 'region' or 'base_url'"):
|
|
TurbopufferDatasink(
|
|
namespace="ns",
|
|
api_key="k",
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# 2. Client initialization
|
|
# =============================================================================
|
|
|
|
|
|
class TestClientInitialization:
|
|
"""Tests for Turbopuffer client lazy initialization."""
|
|
|
|
def test_lazy_initialization(self, sink, mock_turbopuffer_module):
|
|
"""Client is created lazily and cached."""
|
|
client1 = sink._get_client()
|
|
client2 = sink._get_client()
|
|
|
|
assert client1 is client2
|
|
mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
region="gcp-us-central1",
|
|
)
|
|
|
|
def test_uses_explicit_region(self, mock_turbopuffer_module):
|
|
"""Client uses the configured region."""
|
|
sink = make_sink(region="custom-region")
|
|
sink._get_client()
|
|
|
|
mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
region="custom-region",
|
|
)
|
|
|
|
def test_uses_base_url(self, mock_turbopuffer_module):
|
|
"""Client uses base_url when region is not provided."""
|
|
sink = make_sink(
|
|
region=None,
|
|
base_url="https://gcp-us-central1.turbopuffer.com",
|
|
)
|
|
sink._get_client()
|
|
|
|
mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
base_url="https://gcp-us-central1.turbopuffer.com",
|
|
)
|
|
|
|
def test_base_url_does_not_pass_region(self, mock_turbopuffer_module):
|
|
"""When base_url is used, region is not passed to the client."""
|
|
sink = make_sink(
|
|
region=None,
|
|
base_url="https://custom.turbopuffer.com",
|
|
)
|
|
sink._get_client()
|
|
|
|
call_kwargs = mock_turbopuffer_module.Turbopuffer.call_args[1]
|
|
assert "region" not in call_kwargs
|
|
assert call_kwargs["base_url"] == "https://custom.turbopuffer.com"
|
|
|
|
def test_region_does_not_pass_base_url(self, mock_turbopuffer_module):
|
|
"""When region is used, base_url is not passed to the client."""
|
|
sink = make_sink(region="gcp-us-central1")
|
|
sink._get_client()
|
|
|
|
call_kwargs = mock_turbopuffer_module.Turbopuffer.call_args[1]
|
|
assert "base_url" not in call_kwargs
|
|
assert call_kwargs["region"] == "gcp-us-central1"
|
|
|
|
|
|
# =============================================================================
|
|
# 3. Arrow table preparation
|
|
# =============================================================================
|
|
|
|
|
|
class TestArrowTablePreparation:
|
|
"""Tests for _prepare_arrow_table."""
|
|
|
|
def test_renames_columns_and_filters_null_ids(self):
|
|
"""Custom columns are renamed and null IDs filtered."""
|
|
table = pa.table(
|
|
{
|
|
"doc_id": [1, 2, None],
|
|
"emb": [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]],
|
|
}
|
|
)
|
|
sink = make_sink(id_column="doc_id", vector_column="emb")
|
|
|
|
prepared = sink._prepare_arrow_table(table)
|
|
|
|
# Null ID row filtered, columns renamed to id/vector
|
|
expected = pa.table(
|
|
{
|
|
"id": [1, 2],
|
|
"vector": [[0.1, 0.2], [0.3, 0.4]],
|
|
}
|
|
)
|
|
assert prepared.equals(expected)
|
|
|
|
def test_missing_id_column_raises(self):
|
|
"""Missing custom ID column raises ValueError."""
|
|
table = pa.table({"other": [1, 2, 3]})
|
|
sink = make_sink(id_column="doc_id")
|
|
|
|
with pytest.raises(ValueError):
|
|
sink._prepare_arrow_table(table)
|
|
|
|
def test_missing_vector_column_raises(self):
|
|
"""Missing vector column raises ValueError."""
|
|
table = pa.table({"id": [1, 2, 3]})
|
|
sink = make_sink(vector_column="embedding")
|
|
|
|
with pytest.raises(ValueError, match="Vector column 'embedding' not found"):
|
|
sink._prepare_arrow_table(table)
|
|
|
|
@pytest.mark.parametrize(
|
|
"existing_col,custom_col,expected_match",
|
|
[
|
|
("id", "doc_id", "already has.*'id' column"),
|
|
("vector", "emb", "already has.*'vector' column"),
|
|
],
|
|
ids=["id_conflict", "vector_conflict"],
|
|
)
|
|
def test_conflicting_column_names_raise(
|
|
self, existing_col, custom_col, expected_match
|
|
):
|
|
"""Raise if table already has target column name."""
|
|
if existing_col == "id":
|
|
table = pa.table(
|
|
{"id": [1, 2], "doc_id": [10, 20], "vector": [[0.1], [0.2]]}
|
|
)
|
|
sink = make_sink(id_column="doc_id")
|
|
else:
|
|
table = pa.table(
|
|
{"id": [1, 2], "vector": [[0.1], [0.2]], "emb": [[0.3], [0.4]]}
|
|
)
|
|
sink = make_sink(vector_column="emb")
|
|
|
|
with pytest.raises(ValueError, match=expected_match):
|
|
sink._prepare_arrow_table(table)
|
|
|
|
|
|
# =============================================================================
|
|
# 4. Single-namespace batching
|
|
# =============================================================================
|
|
|
|
|
|
class TestSingleNamespaceBatching:
|
|
"""Tests for write batching behavior."""
|
|
|
|
def test_batches_by_batch_size(self, mock_client):
|
|
"""Large tables are split into batches."""
|
|
num_rows = 25
|
|
table = pa.table(
|
|
{
|
|
"id": list(range(num_rows)),
|
|
"vector": [[float(i)] for i in range(num_rows)],
|
|
}
|
|
)
|
|
sink = make_sink(batch_size=10)
|
|
batch_sizes: List[int] = []
|
|
|
|
def track_batch(ns, batch, namespace_name=None):
|
|
# batch is a RecordBatch, get its row count
|
|
batch_sizes.append(batch.num_rows)
|
|
|
|
with patch.object(sink, "_get_client", return_value=mock_client):
|
|
with patch.object(sink, "_write_batch_with_retry", side_effect=track_batch):
|
|
sink.write([table], ctx=None)
|
|
|
|
assert batch_sizes == [10, 10, 5]
|
|
|
|
def test_skips_empty_blocks(self, sink):
|
|
"""Empty blocks don't trigger namespace writes."""
|
|
empty_table = pa.table({"id": [], "vector": []})
|
|
|
|
with patch.object(sink, "_get_client") as mock_get_client:
|
|
with patch.object(sink, "_write_batch_with_retry") as mock_write:
|
|
mock_get_client.return_value = MagicMock()
|
|
sink.write([empty_table], ctx=None)
|
|
|
|
mock_write.assert_not_called()
|
|
|
|
|
|
# =============================================================================
|
|
# 5. Transform to Turbopuffer format
|
|
# =============================================================================
|
|
|
|
|
|
class TestTransformToTurbopufferFormat:
|
|
"""Tests for _transform_to_turbopuffer_format."""
|
|
|
|
def test_requires_id_column(self, sink):
|
|
"""Table must have 'id' column."""
|
|
table = pa.table({"col": [1, 2, 3]})
|
|
|
|
with pytest.raises(ValueError):
|
|
sink._transform_to_turbopuffer_format(table)
|
|
|
|
def test_converts_uuid_bytes_to_native_uuid(self, sink):
|
|
"""16-byte binary IDs become native uuid.UUID objects.
|
|
|
|
Per Turbopuffer performance docs, native UUIDs (16 bytes) are more
|
|
efficient than string UUIDs (36 bytes).
|
|
"""
|
|
u = uuid.uuid4()
|
|
|
|
# ID column must be binary(16) for UUID conversion
|
|
table = pa.table(
|
|
{
|
|
"id": pa.array([u.bytes], type=pa.binary(16)),
|
|
"vector": [[0.1, 0.2]],
|
|
}
|
|
)
|
|
|
|
columns = sink._transform_to_turbopuffer_format(table)
|
|
|
|
expected = {
|
|
"id": [u], # Native uuid.UUID, not bytes
|
|
"vector": [[0.1, 0.2]],
|
|
}
|
|
assert columns == expected
|
|
assert isinstance(columns["id"][0], uuid.UUID)
|
|
|
|
|
|
# =============================================================================
|
|
# 6. Retry logic
|
|
# =============================================================================
|
|
|
|
|
|
class TestRetryLogic:
|
|
"""Tests for _write_batch_with_retry."""
|
|
|
|
@pytest.fixture
|
|
def sample_batch(self):
|
|
"""A simple batch for retry tests."""
|
|
return pa.table({"id": [1], "vector": [[0.1]]})
|
|
|
|
def test_success_first_try(self, sink, sample_batch):
|
|
"""Successful write on first attempt."""
|
|
namespace = MagicMock()
|
|
|
|
sink._write_batch_with_retry(namespace, sample_batch)
|
|
|
|
namespace.write.assert_called_once_with(
|
|
upsert_columns={"id": [1], "vector": [[0.1]]},
|
|
schema=None,
|
|
distance_metric="cosine_distance",
|
|
)
|
|
|
|
def test_retries_then_succeeds(self, sink, sample_batch, monkeypatch):
|
|
"""Transient failures are retried."""
|
|
monkeypatch.setattr(time, "sleep", lambda _: None)
|
|
namespace = MagicMock()
|
|
attempts = {"count": 0}
|
|
|
|
def flaky_write(*args, **kwargs):
|
|
attempts["count"] += 1
|
|
if attempts["count"] < 3:
|
|
raise RuntimeError("temporary error")
|
|
|
|
namespace.write.side_effect = flaky_write
|
|
|
|
sink._write_batch_with_retry(namespace, sample_batch)
|
|
|
|
assert attempts["count"] == 3
|
|
|
|
def test_exhausts_retries_and_raises(self, sink, sample_batch, monkeypatch):
|
|
"""Persistent failures exhaust retries and raise."""
|
|
monkeypatch.setattr(time, "sleep", lambda _: None)
|
|
namespace = MagicMock()
|
|
namespace.write.side_effect = RuntimeError("persistent error")
|
|
|
|
with pytest.raises(RuntimeError, match="persistent error"):
|
|
sink._write_batch_with_retry(namespace, sample_batch)
|
|
|
|
assert namespace.write.call_count == 5 # max_attempts=5
|
|
|
|
@pytest.mark.parametrize(
|
|
"schema,distance_metric",
|
|
[
|
|
({"field": "value"}, "cosine_distance"),
|
|
(None, "euclidean_squared"),
|
|
({"type": "string"}, "euclidean_squared"),
|
|
],
|
|
ids=["with_schema", "alt_metric", "both"],
|
|
)
|
|
def test_configurable_options(self, schema, distance_metric):
|
|
"""Schema and distance_metric are passed to write."""
|
|
sink = make_sink(schema=schema, distance_metric=distance_metric)
|
|
namespace = MagicMock()
|
|
batch = pa.table({"id": [1], "vector": [[0.1]]})
|
|
|
|
sink._write_batch_with_retry(namespace, batch)
|
|
|
|
namespace.write.assert_called_once_with(
|
|
upsert_columns={"id": [1], "vector": [[0.1]]},
|
|
schema=schema,
|
|
distance_metric=distance_metric,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# 7. End-to-end write orchestration
|
|
# =============================================================================
|
|
|
|
|
|
class TestWriteOrchestration:
|
|
"""Tests for top-level write() method."""
|
|
|
|
def test_write_multiple_blocks(self, sink):
|
|
"""Multiple blocks are processed and written."""
|
|
blocks = [
|
|
pa.table({"id": [1, 2], "vector": [[1.0], [2.0]]}),
|
|
pa.table({"id": [3], "vector": [[3.0]]}),
|
|
]
|
|
write_calls = []
|
|
|
|
def track_write(ns, batch, namespace_name=None):
|
|
write_calls.append(batch.num_rows)
|
|
|
|
with patch.object(sink, "_get_client") as mock_get_client:
|
|
mock_client = MagicMock()
|
|
mock_get_client.return_value = mock_client
|
|
with patch.object(sink, "_write_batch_with_retry", side_effect=track_write):
|
|
sink.write(blocks, ctx=None)
|
|
|
|
# Two blocks written
|
|
assert len(write_calls) == 2
|
|
assert write_calls == [2, 1]
|
|
|
|
# Namespace accessed with correct name
|
|
mock_client.namespace.assert_called_with("default_ns")
|
|
|
|
|
|
# =============================================================================
|
|
# 8. Streaming behavior (memory efficiency)
|
|
# =============================================================================
|
|
|
|
|
|
class TestStreamingBehavior:
|
|
"""Tests for memory-efficient streaming writes."""
|
|
|
|
def test_processes_blocks_independently(self, sink):
|
|
"""Each block is processed and written separately."""
|
|
blocks = [pa.table({"id": [i], "vector": [[float(i)]]}) for i in range(5)]
|
|
write_counts = []
|
|
|
|
def track_write(ns, batch, namespace_name=None):
|
|
write_counts.append(batch.num_rows)
|
|
|
|
with patch.object(sink, "_get_client", return_value=MagicMock()):
|
|
with patch.object(sink, "_write_batch_with_retry", side_effect=track_write):
|
|
sink.write(blocks, ctx=None)
|
|
|
|
# 5 blocks → 5 writes of 1 row each
|
|
assert len(write_counts) == 5
|
|
assert all(c == 1 for c in write_counts)
|
|
|
|
|
|
# =============================================================================
|
|
# 9. Multi-namespace writes
|
|
# =============================================================================
|
|
|
|
|
|
class TestMultiNamespaceWrites:
|
|
"""Tests for namespace_column-driven multi-namespace writes."""
|
|
|
|
def test_routes_rows_to_correct_namespaces(self):
|
|
"""Rows are grouped by namespace_column and written to the right ns."""
|
|
sink = make_sink(namespace=None, namespace_column="tenant")
|
|
table = pa.table(
|
|
{
|
|
"tenant": ["ns_a", "ns_b", "ns_a", "ns_b"],
|
|
"id": [1, 2, 3, 4],
|
|
"vector": [[0.1], [0.2], [0.3], [0.4]],
|
|
}
|
|
)
|
|
|
|
writes = {} # namespace_name -> list of row counts
|
|
|
|
def track_write(ns, batch, namespace_name=None):
|
|
writes.setdefault(namespace_name, []).append(batch.num_rows)
|
|
|
|
mock_client = MagicMock()
|
|
mock_client.namespace.return_value = MagicMock()
|
|
with patch.object(sink, "_get_client", return_value=mock_client):
|
|
with patch.object(sink, "_write_batch_with_retry", side_effect=track_write):
|
|
sink.write([table], ctx=None)
|
|
|
|
assert "ns_a" in writes
|
|
assert "ns_b" in writes
|
|
assert sum(writes["ns_a"]) == 2
|
|
assert sum(writes["ns_b"]) == 2
|
|
|
|
def test_drops_namespace_column_before_writing(self):
|
|
"""The namespace column is not included in the written data."""
|
|
sink = make_sink(namespace=None, namespace_column="tenant")
|
|
table = pa.table(
|
|
{
|
|
"tenant": ["ns_a"],
|
|
"id": [1],
|
|
"vector": [[0.1]],
|
|
}
|
|
)
|
|
|
|
written_batches = []
|
|
|
|
def capture_batch(ns, batch, namespace_name=None):
|
|
written_batches.append(batch)
|
|
|
|
mock_client = MagicMock()
|
|
mock_client.namespace.return_value = MagicMock()
|
|
with patch.object(sink, "_get_client", return_value=mock_client):
|
|
with patch.object(
|
|
sink, "_write_batch_with_retry", side_effect=capture_batch
|
|
):
|
|
sink.write([table], ctx=None)
|
|
|
|
assert len(written_batches) == 1
|
|
assert "tenant" not in written_batches[0].column_names
|
|
assert "id" in written_batches[0].column_names
|
|
|
|
def test_missing_namespace_column_raises(self):
|
|
"""Missing namespace column in data raises ValueError."""
|
|
sink = make_sink(namespace=None, namespace_column="tenant")
|
|
table = pa.table(
|
|
{
|
|
"id": [1],
|
|
"vector": [[0.1]],
|
|
}
|
|
)
|
|
|
|
mock_client = MagicMock()
|
|
with patch.object(sink, "_get_client", return_value=mock_client):
|
|
with pytest.raises(ValueError, match="Namespace column.*not found"):
|
|
sink.write([table], ctx=None)
|
|
|
|
def test_null_namespace_values_raise(self):
|
|
"""Null values in namespace column raise ValueError."""
|
|
sink = make_sink(namespace=None, namespace_column="tenant")
|
|
table = pa.table(
|
|
{
|
|
"tenant": ["ns_a", None],
|
|
"id": [1, 2],
|
|
"vector": [[0.1], [0.2]],
|
|
}
|
|
)
|
|
|
|
mock_client = MagicMock()
|
|
with patch.object(sink, "_get_client", return_value=mock_client):
|
|
with pytest.raises(ValueError, match="contains null values"):
|
|
sink.write([table], ctx=None)
|
|
|
|
def test_skips_empty_blocks_in_multi_namespace(self):
|
|
"""Empty blocks are skipped in multi-namespace mode."""
|
|
sink = make_sink(namespace=None, namespace_column="tenant")
|
|
empty_table = pa.table(
|
|
{
|
|
"tenant": pa.array([], type=pa.string()),
|
|
"id": pa.array([], type=pa.int64()),
|
|
"vector": pa.array([], type=pa.list_(pa.float64())),
|
|
}
|
|
)
|
|
|
|
mock_client = MagicMock()
|
|
with patch.object(sink, "_get_client", return_value=mock_client):
|
|
with patch.object(sink, "_write_batch_with_retry") as mock_write:
|
|
sink.write([empty_table], ctx=None)
|
|
|
|
mock_write.assert_not_called()
|
|
|
|
|
|
# =============================================================================
|
|
# 10. Serialization behavior
|
|
# =============================================================================
|
|
|
|
|
|
class TestSerialization:
|
|
"""Tests for pickle serialization support."""
|
|
|
|
def test_preserves_configuration(self, sink, mock_turbopuffer_module):
|
|
"""Configuration is preserved after pickle round-trip."""
|
|
pickled = pickle.dumps(sink)
|
|
unpickled = pickle.loads(pickled)
|
|
|
|
assert unpickled.namespace == sink.namespace
|
|
assert unpickled.namespace_column == sink.namespace_column
|
|
assert unpickled.api_key == sink.api_key
|
|
assert unpickled.region == sink.region
|
|
assert unpickled.base_url == sink.base_url
|
|
assert unpickled.batch_size == sink.batch_size
|
|
assert unpickled._client is None
|
|
|
|
# Lazy initialization works after unpickling
|
|
client = unpickled._get_client()
|
|
assert client is not None
|
|
mock_turbopuffer_module.Turbopuffer.assert_called()
|
|
|
|
def test_preserves_namespace_column_configuration(self, mock_turbopuffer_module):
|
|
"""namespace_column configuration survives pickle round-trip."""
|
|
sink = make_sink(namespace=None, namespace_column="tenant")
|
|
pickled = pickle.dumps(sink)
|
|
unpickled = pickle.loads(pickled)
|
|
|
|
assert unpickled.namespace is None
|
|
assert unpickled.namespace_column == "tenant"
|
|
assert unpickled._client is None
|
|
|
|
def test_preserves_base_url_configuration(self, mock_turbopuffer_module):
|
|
"""base_url configuration survives pickle round-trip."""
|
|
sink = make_sink(
|
|
region=None,
|
|
base_url="https://gcp-us-central1.turbopuffer.com",
|
|
)
|
|
pickled = pickle.dumps(sink)
|
|
unpickled = pickle.loads(pickled)
|
|
|
|
assert unpickled.region is None
|
|
assert unpickled.base_url == "https://gcp-us-central1.turbopuffer.com"
|
|
assert unpickled._client is None
|
|
|
|
# Lazy initialization works and uses base_url
|
|
unpickled._get_client()
|
|
mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
|
|
api_key="test-api-key",
|
|
base_url="https://gcp-us-central1.turbopuffer.com",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|