From bf29caeed001f4f1faa15d3485417c6974c95f53 Mon Sep 17 00:00:00 2001 From: Ben Poppy Date: Thu, 16 Dec 2021 17:37:03 +0000 Subject: [PATCH 1/3] feat: add select function for collection --- mockfirestore/collection.py | 4 ++++ mockfirestore/query.py | 10 +++++++++- tests/test_collection_reference.py | 8 ++++++++ 3 files changed, 21 insertions(+), 1 deletion(-) diff --git a/mockfirestore/collection.py b/mockfirestore/collection.py index 7302ab8..41976f3 100644 --- a/mockfirestore/collection.py +++ b/mockfirestore/collection.py @@ -57,6 +57,10 @@ def offset(self, offset: int) -> Query: query = Query(self, offset=offset) return query + def select(self, fields: List[str]) -> Query: + query = Query(self, select=fields) + return query + def start_at(self, document_fields_or_snapshot: Union[dict, DocumentSnapshot]) -> Query: query = Query(self, start_at=(document_fields_or_snapshot, True)) return query diff --git a/mockfirestore/query.py b/mockfirestore/query.py index 4761a92..f774b7a 100644 --- a/mockfirestore/query.py +++ b/mockfirestore/query.py @@ -8,7 +8,7 @@ class Query: def __init__(self, parent: 'CollectionReference', projection=None, - field_filters=(), orders=(), limit=None, offset=None, + field_filters=(), orders=(), limit=None, offset=None, select=None, start_at=None, end_at=None, all_descendants=False) -> None: self.parent = parent self.projection = projection @@ -16,6 +16,7 @@ def __init__(self, parent: 'CollectionReference', projection=None, self.orders = list(orders) self._limit = limit self._offset = offset + self._select = select self._start_at = start_at self._end_at = end_at self.all_descendants = all_descendants @@ -50,6 +51,9 @@ def stream(self, transaction=None) -> Iterator[DocumentSnapshot]: if self._limit: doc_snapshots = islice(doc_snapshots, self._limit) + if self._select: + doc_snapshots = [DocumentSnapshot(x.reference, {k: v for k, v in x.to_dict().items() if k in self._select}) for x in doc_snapshots] + return iter(doc_snapshots) def get(self) -> Iterator[DocumentSnapshot]: @@ -77,6 +81,10 @@ def offset(self, offset_amount: int) -> 'Query': self._offset = offset_amount return self + def select(self, select_fields: List[str]) -> 'Query': + self._select = select_fields + return self + def start_at(self, document_fields_or_snapshot: Union[dict, DocumentSnapshot]) -> 'Query': self._start_at = (document_fields_or_snapshot, True) return self diff --git a/tests/test_collection_reference.py b/tests/test_collection_reference.py index 59397be..645aa2e 100644 --- a/tests/test_collection_reference.py +++ b/tests/test_collection_reference.py @@ -243,6 +243,14 @@ def test_collection_orderby_offset(self): self.assertEqual({'id': 3}, docs[1].to_dict()) self.assertEqual(2, len(docs)) + def test_collection_select(self): + fs = MockFirestore() + fs._data = {'foo': {'first': {'id': 1, 'value': 'one'}, 'second': {'id': 2, 'value': 'two'}, 'third': {'id': 3, 'value': 'three'}}} + docs = list(fs.collection('foo').select(['id']).stream()) + self.assertEquals({'id': 1}, docs[0].to_dict()) + self.assertEquals({'id': 2}, docs[1].to_dict()) + self.assertEquals({'id': 3}, docs[2].to_dict()) + def test_collection_start_at(self): fs = MockFirestore() fs._data = {'foo': { From e422000079c47d68143460e963b42bca4f6c012b Mon Sep 17 00:00:00 2001 From: Ben Poppy Date: Tue, 25 Jan 2022 13:12:34 +0000 Subject: [PATCH 2/3] refactor: move select under projection in query --- mockfirestore/collection.py | 4 ++-- mockfirestore/query.py | 17 +++++++++++------ 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/mockfirestore/collection.py b/mockfirestore/collection.py index 41976f3..b4f54ad 100644 --- a/mockfirestore/collection.py +++ b/mockfirestore/collection.py @@ -57,8 +57,8 @@ def offset(self, offset: int) -> Query: query = Query(self, offset=offset) return query - def select(self, fields: List[str]) -> Query: - query = Query(self, select=fields) + def select(self, field_paths: List[str]) -> Query: + query = Query(self, projection=field_paths) return query def start_at(self, document_fields_or_snapshot: Union[dict, DocumentSnapshot]) -> Query: diff --git a/mockfirestore/query.py b/mockfirestore/query.py index f774b7a..5aaa2fb 100644 --- a/mockfirestore/query.py +++ b/mockfirestore/query.py @@ -8,7 +8,7 @@ class Query: def __init__(self, parent: 'CollectionReference', projection=None, - field_filters=(), orders=(), limit=None, offset=None, select=None, + field_filters=(), orders=(), limit=None, offset=None, start_at=None, end_at=None, all_descendants=False) -> None: self.parent = parent self.projection = projection @@ -16,7 +16,6 @@ def __init__(self, parent: 'CollectionReference', projection=None, self.orders = list(orders) self._limit = limit self._offset = offset - self._select = select self._start_at = start_at self._end_at = end_at self.all_descendants = all_descendants @@ -51,8 +50,14 @@ def stream(self, transaction=None) -> Iterator[DocumentSnapshot]: if self._limit: doc_snapshots = islice(doc_snapshots, self._limit) - if self._select: - doc_snapshots = [DocumentSnapshot(x.reference, {k: v for k, v in x.to_dict().items() if k in self._select}) for x in doc_snapshots] + if self.projection: + doc_snapshots = [ + DocumentSnapshot( + x.reference, + {k: v for k, v in x.to_dict().items() if k in self.projection} + ) + for x in doc_snapshots + ] return iter(doc_snapshots) @@ -81,8 +86,8 @@ def offset(self, offset_amount: int) -> 'Query': self._offset = offset_amount return self - def select(self, select_fields: List[str]) -> 'Query': - self._select = select_fields + def select(self, field_paths: List[str]) -> 'Query': + self.projection = field_paths return self def start_at(self, document_fields_or_snapshot: Union[dict, DocumentSnapshot]) -> 'Query': From 90e134dc0ce9c505710de68f76832c3ed493173a Mon Sep 17 00:00:00 2001 From: Ben Poppy Date: Tue, 25 Jan 2022 13:12:53 +0000 Subject: [PATCH 3/3] test: add additional tests for select --- tests/test_collection_reference.py | 36 ++++++++++++++++++++++++++---- 1 file changed, 32 insertions(+), 4 deletions(-) diff --git a/tests/test_collection_reference.py b/tests/test_collection_reference.py index 645aa2e..0bdc75d 100644 --- a/tests/test_collection_reference.py +++ b/tests/test_collection_reference.py @@ -245,11 +245,39 @@ def test_collection_orderby_offset(self): def test_collection_select(self): fs = MockFirestore() - fs._data = {'foo': {'first': {'id': 1, 'value': 'one'}, 'second': {'id': 2, 'value': 'two'}, 'third': {'id': 3, 'value': 'three'}}} + fs._data = {'foo': { + 'first': {'id': 1, 'value': 'one'}, + 'second': {'id': 2, 'value': 'two'}, + 'third': {'id': 3, 'value': 'three'} + }} docs = list(fs.collection('foo').select(['id']).stream()) - self.assertEquals({'id': 1}, docs[0].to_dict()) - self.assertEquals({'id': 2}, docs[1].to_dict()) - self.assertEquals({'id': 3}, docs[2].to_dict()) + self.assertEqual({'id': 1}, docs[0].to_dict()) + self.assertEqual({'id': 2}, docs[1].to_dict()) + self.assertEqual({'id': 3}, docs[2].to_dict()) + + def test_collection_select_missing_fields(self): + fs = MockFirestore() + fs._data = {'foo': { + 'first': {'id': 1, 'value': 'one'}, + 'second': {'id': 2, 'value': 'two'}, + 'third': {'id': 3, 'value': 'three'} + }} + docs = list(fs.collection("foo").select(["name"]).stream()) + self.assertEqual({}, docs[0].to_dict()) + self.assertEqual({}, docs[1].to_dict()) + self.assertEqual({}, docs[2].to_dict()) + + def test_collection_select_one_missing_field(self): + fs = MockFirestore() + fs._data = {'foo': { + 'first': {'id': 1, 'value': 'one'}, + 'second': {'id': 2, 'value': 'two'}, + 'third': {'id': 3, 'value': 'three'} + }} + docs = list(fs.collection("foo").select(["id", "name"]).stream()) + self.assertEqual({"id": 1}, docs[0].to_dict()) + self.assertEqual({"id": 2}, docs[1].to_dict()) + self.assertEqual({"id": 3}, docs[2].to_dict()) def test_collection_start_at(self): fs = MockFirestore()