Fix for python 2.x

This commit is contained in:
Ryan Leckey
2014-12-12 11:19:12 -08:00
parent 15dcbb1054
commit d21e400a8b
2 changed files with 13 additions and 10 deletions

View File

@@ -28,7 +28,7 @@ class EncryptionDecryptionBaseEngine(object):
def _update_key(self, key):
if isinstance(key, six.string_types):
key = six.b(key)
key = key.encode()
digest = hashes.Hash(hashes.SHA256(), backend=default_backend())
digest.update(key)
engine_key = digest.finalize()

View File

@@ -139,7 +139,7 @@ class EncryptedTypeTestCase(TestCase):
class Team(self.Base):
__tablename__ = 'team'
id = sa.Column(sa.Integer, primary_key=True)
key = sa.Column(sa.Unicode(50))
key = sa.Column(sa.String(50))
name = sa.Column(EncryptedType(
sa.Unicode,
lambda: self._team_key,
@@ -181,18 +181,17 @@ class EncryptedTypeTestCase(TestCase):
def test_lookup_key(self):
# Add teams
self._team_key = 'one'
team = self.Team(key=self._team_key, name='One')
team = self.Team(key=self._team_key, name=u'One')
self.session.add(team)
self.session.flush()
self.session.commit()
team_1_id = team.id
self._team_key = 'two'
team = self.Team(key=self._team_key, name='Two')
team = self.Team(key=self._team_key)
team.name = u'Two'
self.session.add(team)
self.session.flush()
team_2_id = team.id
self.session.commit()
team_2_id = team.id
# Lookup teams
self._team_key = self.session.query(self.Team.key).filter_by(
@@ -201,22 +200,26 @@ class EncryptedTypeTestCase(TestCase):
team = self.session.query(self.Team).get(team_1_id)
assert team.name == 'One'
assert team.name == u'One'
with pytest.raises(Exception):
self.session.query(self.Team).get(team_2_id)
self.session.expunge_all()
self._team_key = self.session.query(self.Team.key).filter_by(
id=team_2_id
).one()[0]
team = self.session.query(self.Team).get(team_2_id)
assert team.name == 'Two'
assert team.name == u'Two'
with pytest.raises(Exception):
self.session.query(self.Team).get(team_1_id)
self.session.expunge_all()
# Remove teams
self.session.query(self.Team).delete()
self.session.commit()