codekingpro/portable-devtools

1"""Tests for the Count API endpoint with read_level support."""
2
3from typing import Tuple
4from uuid import uuid4
5
6import pytest
7
8from chromadb.api import ClientAPI
9from chromadb.api.models.Collection import Collection
10from chromadb.api.types import ReadLevel
11from chromadb.test.conftest import (
12    ClientFactories,
13    is_spann_disabled_mode,
14    skip_reason_spann_disabled,
15)
16
17
18def _create_test_collection(
19    client_factories: ClientFactories,
20) -> Tuple[Collection, ClientAPI]:
21    """Create a test collection with some data."""
22    client = client_factories.create_client_from_system()
23    client.reset()
24
25    collection_name = f"count_api_test_{uuid4().hex}"
26    collection = client.get_or_create_collection(name=collection_name)
27
28    return collection, client
29
30
31@pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
32def test_count_with_read_level_index_and_wal(
33    client_factories: ClientFactories,
34) -> None:
35    """Test count with ReadLevel.INDEX_AND_WAL (default) returns correct count."""
36    collection, _ = _create_test_collection(client_factories)
37
38    collection.add(
39        ids=["doc1", "doc2", "doc3"],
40        documents=["apple fruit", "banana fruit", "car vehicle"],
41        embeddings=[[0.1, 0.2, 0.3, 0.4], [0.2, 0.3, 0.4, 0.5], [0.9, 0.8, 0.7, 0.6]],
42    )
43
44    count = collection.count(read_level=ReadLevel.INDEX_AND_WAL)
45    assert count == 3
46
47
48@pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
49def test_count_with_read_level_index_only(
50    client_factories: ClientFactories,
51) -> None:
52    """Test count with ReadLevel.INDEX_ONLY returns a valid count."""
53    collection, _ = _create_test_collection(client_factories)
54
55    collection.add(
56        ids=["doc1", "doc2", "doc3"],
57        documents=["apple fruit", "banana fruit", "car vehicle"],
58        embeddings=[[0.1, 0.2, 0.3, 0.4], [0.2, 0.3, 0.4, 0.5], [0.9, 0.8, 0.7, 0.6]],
59    )
60
61    # Count with INDEX_ONLY skips the WAL.
62    # Result may be less than 3 if data hasn't been compacted yet.
63    count = collection.count(read_level=ReadLevel.INDEX_ONLY)
64    assert isinstance(count, int)
65    assert count >= 0
66
67
68@pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
69def test_count_with_read_level_index_and_bounded_wal(
70    client_factories: ClientFactories,
71) -> None:
72    """Test count with ReadLevel.INDEX_AND_BOUNDED_WAL returns a valid count."""
73    collection, _ = _create_test_collection(client_factories)
74
75    collection.add(
76        ids=["doc1", "doc2", "doc3"],
77        documents=["apple fruit", "banana fruit", "car vehicle"],
78        embeddings=[[0.1, 0.2, 0.3, 0.4], [0.2, 0.3, 0.4, 0.5], [0.9, 0.8, 0.7, 0.6]],
79    )
80
81    # Count with INDEX_AND_BOUNDED_WAL reads up to a server-configured number of WAL entries.
82    # Result depends on the configured limit and WAL size.
83    count = collection.count(read_level=ReadLevel.INDEX_AND_BOUNDED_WAL)
84    assert isinstance(count, int)
85    assert count >= 0
86
87
88@pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
89def test_count_default_read_level(
90    client_factories: ClientFactories,
91) -> None:
92    """Test count without explicit read_level uses default (INDEX_AND_WAL)."""
93    collection, _ = _create_test_collection(client_factories)
94
95    collection.add(
96        ids=["doc1", "doc2"],
97        documents=["hello world", "goodbye world"],
98        embeddings=[[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]],
99    )
100
101    count = collection.count()
102    assert count == 2
103