From 4714e013ba2e8559b4eda994a03febe55c49204e Mon Sep 17 00:00:00 2001 From: Konsta Vesterinen Date: Wed, 14 May 2014 15:47:22 +0300 Subject: [PATCH] Refactor tests --- tests/functions/test_get_columns.py | 11 +++++++---- tests/functions/test_get_primary_keys.py | 10 +++++++--- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/tests/functions/test_get_columns.py b/tests/functions/test_get_columns.py index 18dd57c..1ed3e3a 100644 --- a/tests/functions/test_get_columns.py +++ b/tests/functions/test_get_columns.py @@ -1,11 +1,14 @@ import sqlalchemy as sa +from sqlalchemy.ext.declarative import declarative_base + from sqlalchemy_utils import get_columns -from tests import TestCase -class TestGetColumns(TestCase): - def create_models(self): - class Building(self.Base): +class TestGetColumns(object): + def setup_method(self, method): + Base = declarative_base() + + class Building(Base): __tablename__ = 'building' id = sa.Column('_id', sa.Integer, primary_key=True) name = sa.Column('_name', sa.Unicode(255)) diff --git a/tests/functions/test_get_primary_keys.py b/tests/functions/test_get_primary_keys.py index c5fadab..09bfa29 100644 --- a/tests/functions/test_get_primary_keys.py +++ b/tests/functions/test_get_primary_keys.py @@ -3,13 +3,17 @@ try: except ImportError: from ordereddict import OrderedDict import sqlalchemy as sa +from sqlalchemy.ext.declarative import declarative_base + from sqlalchemy_utils import get_primary_keys from tests import TestCase -class TestGetPrimaryKeys(TestCase): - def create_models(self): - class Building(self.Base): +class TestGetPrimaryKeys(object): + def setup_method(self, method): + Base = declarative_base() + + class Building(Base): __tablename__ = 'building' id = sa.Column('_id', sa.Integer, primary_key=True) name = sa.Column('_name', sa.Unicode(255))