diff --git a/scenarious/scenario.py b/scenarious/scenario.py index f4a48d0..2d9e0cd 100644 --- a/scenarious/scenario.py +++ b/scenarious/scenario.py @@ -54,7 +54,7 @@ def update(self, source): if isinstance(source, dict): raw = source else: - raw = yaml.load(open(source) if isinstance(source, six.string_types) else source) + raw = yaml.full_load(open(source) if isinstance(source, six.string_types) else source) for entity, value in (raw or {}).items(): objects = [{}] * value if isinstance(value, int) else value @@ -94,7 +94,7 @@ def __getattr__(self, key): type_name = self._get_type_name(key) if self._entity_store.has_type(type_name): - return list(self._entity_store.all(type_name)) + return self._entity_store.all(type_name) else: raise AttributeError("%s doesn't have type '%s'" % (self.__class__.__name__, type_name)) diff --git a/scenarious/store_handler.py b/scenarious/store_handler.py index d8fe2ed..22637bd 100644 --- a/scenarious/store_handler.py +++ b/scenarious/store_handler.py @@ -13,6 +13,52 @@ class EntityStoreException(Exception): pass +class EntityManager(object): + + def __init__(self, type_name, objects=None, aliased_objects=None): + self.type_name = type_name + self.identifiers = [str(id) for id in objects.keys()] + self.objects = list(objects.values()) + if aliased_objects: + self.aliases = list(aliased_objects.keys()) + self.aliased_objects = list(aliased_objects.values()) + + def __repr__(self): + return str(self.objects) + + def __getitem__(self, key): + try: + item = self.objects[key] + return item + except: + raise EntityStoreException( + "{} type name has no object with identifier: {}".format(self.type_name, key) + ) + + def __len__(self): + return len(self.objects) + + def __contains__(self, item): + return item in self.objects + + def __getattr__(self, name): + attr_parts = name.split('_') + requested_type_name, requested_identifier = attr_parts + + if requested_type_name != self.type_name: + raise EntityStoreException("Invalid type name: {}".format(requested_type_name)) + + if requested_identifier not in self.identifiers and requested_identifier not in self.aliases: + raise EntityStoreException("Object not found: {}".format(name)) + + try: + object_index = self.identifiers.index(requested_identifier) + return self.objects[object_index] + except ValueError: + object_index = self.aliases.index(requested_identifier) + return self.aliased_objects[object_index] + + class EntityStore(object): ID = 'id' @@ -99,4 +145,6 @@ def get(self, type_name, ref): return e def all(self, type_name): - return self._objects.get(type_name, {}).values() + objects = self._objects.get(type_name, {}) + aliased_objects = self._aliased_objects.get(type_name, {}) + return EntityManager(type_name, objects=objects, aliased_objects=aliased_objects) diff --git a/tests/test_scenarious.py b/tests/test_scenarious.py index 87ff052..8ac6325 100644 --- a/tests/test_scenarious.py +++ b/tests/test_scenarious.py @@ -8,6 +8,7 @@ from io import StringIO from scenarious.util import DictObject +from scenarious.store_handler import EntityStoreException from scenarious.type_handlers.base import faker, TypeHandlerException from scenarious import Scenario, TypeHandler, ScenariousException @@ -72,11 +73,17 @@ def test_load_type(self): assert 1 == len(s.actors) assert 'test' == s.actors[0].name assert 20 == s.actors[0].age + # Testing attribute access to objects + assert "test" == s.actors.actor_1.name + assert 20 == s.actors.actor_1.age assert s.movies assert 1 == len(s.movies) assert 'test movie' == s.movies[0].title assert 'drama' == s.movies[0].genre + # Testing attribute access to objects + assert 'test movie' == s.movies.movie_1.title + assert 'drama' == s.movies.movie_1.genre def test_custom_ids(self): s = Scenario.load(StringIO(""" @@ -102,6 +109,15 @@ def test_custom_ids(self): assert 'test1' == actor_1.name assert 'test2' == actor_3.name + # Testing direct attribute access to objects with custom ids + actor_1 = s.actors.actor_1 + actor_3 = s.actors.actor_3 + assert 1 == actor_1.id + assert 3 == actor_3.id + + assert 'test1' == actor_1.name + assert 'test2' == actor_3.name + def test_alias_objects(self): s = Scenario.load(StringIO(""" actors: @@ -121,6 +137,11 @@ def test_alias_objects(self): assert 'test1' == s.by_id('actors', 'test1').name assert not s.by_id('actors', 'test2') # id:3 has no alias + # Testing direct attribute access to objects with custom ids + assert 'test1' == s.actors.actor_test1.name + with self.assertRaises(EntityStoreException, msg="Object not found: test2"): + not s.actors.actor_test2 + def test_fail_with_duplicate_alias(self): scene = StringIO(""" actors: @@ -158,6 +179,11 @@ def test_mix_auto_generated_and_custom_ids(self): # should've been forced to relocate because of test33 having _id:1 assert 'test1' == s.by_id('actors', 3).name + # Testing direct attribute access to objects with custom ids + assert 'test3' == s.actors.actor_1.name + assert 'test2' == s.actors.actor_2.name + assert 'test1' == s.actors.actor_3.name + def test_get_object_by_id(self): s = Scenario.load(StringIO(""" actors: