diff --git a/examples/distributed_counter_example.py b/examples/distributed_counter_example.py new file mode 100644 index 0000000..9e1d11c --- /dev/null +++ b/examples/distributed_counter_example.py @@ -0,0 +1,78 @@ +from flask_boiler.context import Context as CTX +from flask_boiler.schema import Schema +from flask_boiler.fields import Integer +from flask_boiler.domain_model import DomainModel +from flask_boiler.view_mediator_dav import ViewMediatorDAV +from flask_boiler.view_model import ViewModel + + +class ShardSchema(Schema): + + count = Integer() + + +class Shard(DomainModel): + + class Meta: + schema_cls = ShardSchema + collection_id = "CounterAShards" + + +class CounterViewSchema(Schema): + + total_count = Integer(dump_only=True) + + +class CounterView(ViewModel): + + class Meta: + schema_cls = CounterViewSchema + + def __init__(self, *args, **kwargs): + doc_ref = CTX.db.collection("counters").document("counter_0") + super().__init__(*args, doc_ref=doc_ref, **kwargs) + self.shards = dict() + + @property + def total_count(self): + return sum(v.count for _, v in self.shards.items()) + + def set_shard(self, sid, shard): + self.shards[sid] = shard + + def get_vm_update_callback(self, dm_cls, *args, **kwargs) : + + if dm_cls == Shard: + def callback(vm: CounterView, dm: Shard): + vm.set_shard(dm.doc_id, dm) + return callback + else: + return super().get_vm_update_callback(dm_cls, *args, **kwargs) + + +class CounterMediator(ViewMediatorDAV): + + def __init__(self, shard_size, *args, **kwargs): + super().__init__(*args, view_model_cls=CounterView, **kwargs) + self.shard_size = shard_size + self.view_model = None + + @classmethod + def notify(cls, obj): + obj.save() + + def start(self): + + struct = dict() + + for i in range(self.shard_size): + doc_id = str(i) + shard = Shard.new(doc_id=doc_id, ) + shard.save() + struct[shard.doc_id] = (Shard, doc_id) + + self.view_model = self.view_model_cls.get( + f_notify=self.notify, + struct_d=struct, + once=False) + diff --git a/tests/test_examples.py b/tests/test_examples.py index 633d5b0..7d62fdd 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1,2 +1,48 @@ +from google.cloud.firestore import Increment + from examples.meeting_room.tests.test_view_model import * from examples.meeting_room.tests.test_dav import * + +from examples import distributed_counter_example as dc +from .fixtures import CTX + + +def test_start(CTX): + """ + TODO: add teardown steps to delete documents generated in firestore + + :param CTX: + :return: + """ + + mediator = dc.CounterMediator( + mutation_cls=None, + shard_size=2 + ) + + mediator.start() + + time.sleep(3) + + doc_ref = CTX.db.collection("counters").document("counter_0") + + assert doc_ref.get().to_dict() == { + "doc_ref": "counters/counter_0", + "obj_type": "CounterView", + "totalCount": 0 + } + + CTX.db.collection("Shard").document("0").set( + document_data={ + "count": Increment(1), + }, + merge=True + ) + + time.sleep(1) + + assert doc_ref.get().to_dict() == { + "doc_ref": "counters/counter_0", + "obj_type": "CounterView", + "totalCount": 1 + }