Skip to content

Commit

Permalink
fix: Adding id in the model for state validation
Browse files Browse the repository at this point in the history
- Problem to solve:  same collection names in multiple tenants/db
  • Loading branch information
tazarov committed Feb 15, 2024
1 parent 6a5439b commit 819685c
Showing 1 changed file with 15 additions and 6 deletions.
21 changes: 15 additions & 6 deletions chromadb/test/property/test_collections.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import math

import pytest
import logging
import hypothesis.strategies as st
Expand Down Expand Up @@ -54,7 +56,7 @@ def create_coll(
metadata=coll.metadata,
embedding_function=coll.embedding_function,
)
self.set_model(coll.name, coll.metadata)
self.set_model(coll.name, coll.metadata,str(coll.id))

assert c.name == coll.name
assert c.metadata == self.model[coll.name]
Expand Down Expand Up @@ -85,7 +87,7 @@ def delete_coll(self, coll: strategies.Collection) -> None:
@rule()
def list_collections(self) -> None:
colls = self.api.list_collections()
assert len(colls) == len(self.model)
assert len(colls) == len([c for c in self.model if not c.startswith("__id__")])
for c in colls:
assert c.name in self.model

Expand Down Expand Up @@ -163,7 +165,7 @@ def get_or_create_coll(
coll.metadata = (
self.model[coll.name] if new_metadata is None else new_metadata
)
self.set_model(coll.name, coll.metadata)
self.set_model(coll.name, coll.metadata,str(coll.id))

# Update API
c = self.api.get_or_create_collection(
Expand All @@ -189,6 +191,9 @@ def modify_coll(
new_metadata: types.Metadata,
new_name: Optional[str],
) -> MultipleResults[strategies.Collection]:
# early exit if a col with name exists but with diff id, possibly in another tenant/db
if coll.name in self.model and f"__id__:{coll.id}" not in self.model:
return multiple()
if coll.name not in self.model:
with pytest.raises(Exception):
c = self.api.get_collection(name=coll.name)
Expand Down Expand Up @@ -218,7 +223,7 @@ def modify_coll(
self.delete_from_model(coll.name)
coll.name = new_name
_name = new_name
self.set_model(_name, _metadata) # type: ignore
self.set_model(_name, _metadata, str(coll.id))

c.modify(metadata=new_metadata, name=new_name)
c = self.api.get_collection(name=coll.name)
Expand All @@ -228,14 +233,18 @@ def modify_coll(
return multiple(coll)

def set_model(
self, name: str, metadata: Optional[types.CollectionMetadata]
self, name: str, metadata: Optional[types.CollectionMetadata], id: Optional[str] = None
) -> None:
model = self.model
model[name] = metadata
if id is not None:
model[f"__id__:{id}"] = metadata

def delete_from_model(self, name: str) -> None:
def delete_from_model(self, name: str, id: Optional[str] = None) -> None:
model = self.model
del model[name]
if id is not None:
del model[f"__id__:{id}"]

@property
def model(self) -> Dict[str, Optional[types.CollectionMetadata]]:
Expand Down

0 comments on commit 819685c

Please sign in to comment.