手动为联接的表继承模型插入测试数据

时间:2019-01-17 08:17:14

标签: python sqlalchemy

我在SQLAlchemy的帮助器类中有一个get_or_create方法,但是会话从未从查询返回任何数据,因此始终尝试创建新行-即使数据库中有相同的数据(然后引发UNIQUE CONSTRANT错误)。

通过create方法本身添加数据确实可以成功完成,并且数据被持久保存到数据库中。在get_or_create()方法中,实例始终为None,因此尝试将数据再次添加到数据库中。

下面的代码是完整应用程序的摘录,应该显示出足以复制该问题的代码。

models.py

from sqlalchemy import Column, Integer, String, ForeignKey, Float, Enum
from sqlalchemy.ext.declarative import declarative_base
import enum

Base = declarative_base()


class DepositStatus(enum.Enum):
    not_yet_paid = 1
    paid = 2
    to_be_refunded = 3
    deposit_kept = 4
    refunded = 5


class RiderStatus(enum.Enum):
    not_yet_started = 1
    active = 2
    finished = 3
    scratched = 4


class RiderCategories(enum.Enum):
    solo = 1
    pair = 2


class RiderGenders(enum.Enum):
    male = 1
    female = 2


class TrackerLocations(Base):
    __tablename__ = 'tracker_locations'
    id = Column('id', Integer, primary_key=True)
    types = Column(String)
    __mapper_args__ = {
        'polymorphic_on': types
    }


class Riders(TrackerLocations):

    __tablename__ = 'riders'
    id = Column(Integer, ForeignKey('tracker_locations.id'), primary_key=True,
                autoincrement=True)
    first_name = Column('first_name', String)
    last_name = Column('last_name', String)
    cap_number = Column('cap_number', String)
    category = Column('category', Enum(RiderCategories))
    status = Column('status', Enum(RiderStatus))
    gender = Column('gender', Enum(RiderGenders))
    deposit_status = Column('deposit_status', Enum(DepositStatus))
    deposit_amount = Column('deposit_amount', Float)
    email = Column('email', String)

    __mapper_args__ = {
        'polymorphic_identity': 'riders'
    }

db.py

from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy import create_engine
from models import Base

class DB(object):
    """
    To be used to interact with SQLAlchemy's ORM.
    """

    def __init__(self, db_uri):
        self.engine = create_engine(db_uri)
        Base.metadata.create_all(self.engine)
        self._session = scoped_session(sessionmaker(bind=self.engine))

    def close(self):
        self.session.close()

    def rollback(self):
        self.session.rollback()

    @property
    def session(self):
        return self._session()

    def get_or_create(self, model, commit=True, **kwargs):
        instance = self.session.query(model).filter_by(**kwargs).first()
        if instance:
            return instance, False
        else:
            instance = self.create(model, commit=commit, **kwargs)
            return instance, True

    def create(self, model, commit=True, **kwargs):
        instance = model(**kwargs)
        self.session.add(instance)
        if commit:
            self.session.commit()
        return instance

tests.py:

import os
import sqlite3
from tempfile import mkdtemp
from unittest import TestCase

from models import Riders
from db import DB

class TestDBInteractions(TestCase):

    @classmethod
    def setUpClass(cls):
        cls.temp_ = mkdtemp()
        cls.temp_db = os.path.join(cls.temp_, 'test_db.db')

    def setUp(self):
        self.db_uri = 'sqlite:///' + self.temp_db
        self.db = DB(self.db_uri)
        self.conn = sqlite3.connect(self.temp_db)
        self.cur = self.conn.cursor()

    def tearDown(self):
        self.db.close()

    def test_get_or_create(self):
        test_data = {
            'id': 1,
            'first_name': 'Bobby',
            'last_name': 'Hill',
            'cap_number': '123',
            'category': 'solo',
            'status': 'finished',
            'gender': 'male',
            'deposit_status': 'paid',
            'deposit_amount': 123.1,
            'email': 'hello@bob.com'
        }
        params = tuple(test_data.values())
        query = "INSERT INTO riders VALUES (?,?,?,?,?,?,?,?,?,?)"
        self.cur.execute(query, params)
        # # sqlite doesn't allow multiple write connections to the same DB, so close this connection
        self.conn.commit()
        self.conn.close()
        self.db1 = DB(db_uri=self.db_uri)
        response = self.db1.get_or_create(Riders, **{'id': 1})
        print(response)

2 个答案:

答案 0 :(得分:1)

使用Joined Table Inheritance,SQLA可以出色地完成工作,让您忘记正在使用两个表。

您的测试设置不会向tracker_locations插入任何数据,就像实例化Riders的实例并通过ORM插入数据一样。

这是SQLA生成的查询的联接部分:

FROM tracker_locations JOIN riders ON tracker_locations.id = riders.id

在测试数据集中,tracker_locations中没有行,因此联接条件tracker_locations.id = riders.id的评估结果为true。

因此,如果您还在tracker_locations表中包含一些测试数据,它将起作用:

tracker_locations_test_data = {'id': 1, 'type': 'riders'}
self.cur.execute(
    'INSERT INTO tracker_locations VALUES (?,?)', 
    tuple(tracker_locations_test_data.values()))

答案 1 :(得分:1)

仅使用sqlite3手动将数据加载到riders表中不足以兑现Riders模型中TrackerLocations模型的“联合继承”。您还必须将正确的数据填充到TrackerLocations中,以进行sqlalchemy才能使用Models中的信息并进行正确的查找。

使用:

def test_get_or_create(self):
    test_data = {
        'id': 1,
        'first_name': 'Bobby',
        'last_name': 'Hill',
        'cap_number': '123',
        'category': 'solo',
        'status': 'finished',
        'gender': 'male',
        'deposit_status': 'paid',
        'deposit_amount': 123.1,
        'email': 'hello@bob.com'
    }
    test_tracker = {
        'id':1,
        'type':'riders'
    }
    params = tuple(test_data.values())
    query = "INSERT INTO riders VALUES (?,?,?,?,?,?,?,?,?,?)"
    self.cur.execute(query, params)
    query = "INSERT INTO tracker_locations VALUES (?,?)"
    params = tuple(test_tracker.values())
    self.cur.execute(query, params)
    # # sqlite doesn't allow multiple write connections to the same DB, so close this connection
    self.conn.commit()
    self.conn.close()

    self.db1 = DB(db_uri=self.db_uri)

    response = self.db1.get_or_create(Riders, **{'id': 1})
    print(response)

用于测试,它应该可以工作。

相关问题