diff --git a/python/tests/test_collection.py b/python/tests/test_collection.py index 7eba2e2..b03d1e3 100644 --- a/python/tests/test_collection.py +++ b/python/tests/test_collection.py @@ -13,6 +13,8 @@ # limitations under the License. from __future__ import annotations +from unittest.mock import MagicMock + import pytest import zvec from zvec import ( @@ -907,6 +909,17 @@ class TestCollectionQuery: assert len(doc.field_names()) == 2 assert set(doc.field_names()) == {"id", "name"} + @pytest.mark.parametrize("topk", [0, -1, None, True]) + def test_collection_query_rejects_invalid_topk(self, topk): + collection = Collection.__new__(Collection) + collection._querier = MagicMock() + collection._obj = MagicMock() + + with pytest.raises(ValueError, match="topk must be a positive integer"): + collection.query(Query(field_name="dense", vector=[0.1]), topk=topk) + + collection._querier.execute.assert_not_called() + def test_collection_query_with_topk( self, collection_with_multiple_docs: Collection ): diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py index 6663a1b..a002d0d 100644 --- a/python/zvec/model/collection.py +++ b/python/zvec/model/collection.py @@ -423,6 +423,7 @@ class Collection: ... output_fields=["title", "url"] ... ) """ + _require_positive_integer(topk, "topk") if vectors is not None: warnings.warn( "The 'vectors' parameter is deprecated and will be removed in a future version. "