diff --git a/sqlalchemy_seed/__init__.py b/sqlalchemy_seed/__init__.py index 9dd3aae..7ccb57b 100644 --- a/sqlalchemy_seed/__init__.py +++ b/sqlalchemy_seed/__init__.py @@ -14,6 +14,8 @@ import os import yaml +from sqlalchemy import MetaData + try: from yaml import CLoader as Loader except ImportError: @@ -90,6 +92,22 @@ def _create_model_instance(fixture): return instances +def _create_table_object_data(fixture, session): + """Create a Table object entry. + + :param fixture: Fixtures + """ + for data in fixture: + if 'table' in data: + module_name, class_name = data['table'].rsplit('.', 1) + importlib.import_module(module_name) + metadata = MetaData() + metadata.reflect(bind=session.get_bind()) + table = metadata.tables[class_name] + insert = table.insert() + session.execute(insert.values(**data['fields'])) + + def load_fixtures(session, fixtures): """Load fixture. @@ -99,6 +117,7 @@ def load_fixtures(session, fixtures): instances = [] for fixture in fixtures: _instances = _create_model_instance(fixture) + _create_table_object_data(fixture, session) for instance in _instances: instances.append(instance) diff --git a/tests/fixtures/picture_categories.yaml b/tests/fixtures/picture_categories.yaml new file mode 100644 index 0000000..5cd7472 --- /dev/null +++ b/tests/fixtures/picture_categories.yaml @@ -0,0 +1,9 @@ +- model: tests.test_sqlalchey_seed.PictureCategory + id: 1 + fields: + name: Category1 +- model: tests.test_sqlalchey_seed.PictureCategory + id: 2 + fields: + name: Category2 + diff --git a/tests/fixtures/picture_category_picture.yaml b/tests/fixtures/picture_category_picture.yaml new file mode 100644 index 0000000..02ec2cc --- /dev/null +++ b/tests/fixtures/picture_category_picture.yaml @@ -0,0 +1,12 @@ +- table: tests.test_sqlalchey_seed.picture_category_picture + id: 1 + fields: + picture_id: 1 + picture_category_id: 1 + +- table: tests.test_sqlalchey_seed.picture_category_picture + id: 2 + fields: + picture_id: 1 + picture_category_id: 2 + diff --git a/tests/test_sqlalchey_seed.py b/tests/test_sqlalchey_seed.py index b1d08fd..cb80d53 100644 --- a/tests/test_sqlalchey_seed.py +++ b/tests/test_sqlalchey_seed.py @@ -13,7 +13,7 @@ import os from unittest import TestCase -from sqlalchemy import Column, create_engine, ForeignKey, Integer, String +from sqlalchemy import Column, create_engine, ForeignKey, Integer, String, Table from sqlalchemy.exc import OperationalError from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import relationship, scoped_session, sessionmaker @@ -75,6 +75,28 @@ def __repr__(self): ) +class PictureCategory(Base): + __tablename__ = 'picture_category' + + id = Column(Integer, primary_key=True) + name = Column(String(120)) + + def __repr__(self): + """Repr.""" + return 'PictureCategory(id={0}, name={1})'.format( + self.id, + self.name, + ) + + +picture_category_picture = Table( + 'picture_category_picture', + Base.metadata, + Column('picture_id', Integer, ForeignKey('pictures.id')), + Column('picture_category_id', Integer, ForeignKey('picture_category.id')), +) + + class TestFixtures(TestCase): path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'fixtures') @@ -105,13 +127,17 @@ def test_load_fixture(self): def test_load_fixtures(self): create_table(Base) fixtures = load_fixture_files( - self.path, ['accounts.yaml', 'pictures.yaml'], + self.path, ['accounts.yaml', 'pictures.yaml', 'picture_categories.yaml', 'picture_category_picture.yaml'], ) load_fixtures(session, fixtures) accounts = session.query(Account).all() self.assertEqual(len(accounts), 2) pictures = session.query(Picture).all() self.assertEqual(len(pictures), 4) + picture_categories = session.query(PictureCategory).all() + self.assertEqual(len(picture_categories), 2) + category_rels = session.query(picture_category_picture).all() + self.assertEqual(len(category_rels), 2) drop_table(Base, session)