Skip to content

Commit

Permalink
Small test improvement
Browse files Browse the repository at this point in the history
  • Loading branch information
ivorbosloper committed Nov 22, 2017
1 parent 5d2efd9 commit b745418
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 5 deletions.
2 changes: 2 additions & 0 deletions tests/tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -463,10 +463,12 @@ def test_migrate_from_login(self):
from django.contrib.sessions.backends.db import SessionStore as DjangoSessionStore
try:
call_command('migrate', 'sessions')
call_command('clearsessions')
user = User.objects.create_user('bouke', '', 'secret')
session = DjangoSessionStore()
session['_auth_user_id'] = user.id
session.save()
self.assertEqual(Session.objects.count(), 0)
self.assertEqual(DjangoSession.objects.count(), 1)
call_command('migratesessions')
new_sessions = list(Session.objects.all())
Expand Down
14 changes: 9 additions & 5 deletions user_sessions/management/commands/migratesessions.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,12 @@ class Command(BaseCommand):
Convert existing (old) sessions to the new session store.
"""
def add_arguments(self, parser):
parser.add_argument('--oldmodel',
dest='oldmodel',
default='django.contrib.sessions.models.Session',
help='Existing session model to migrate to the new UserSessions database table')
parser.add_argument(
'--oldmodel',
dest='oldmodel',
default='django.contrib.sessions.models.Session',
help='Existing session model to migrate to the new UserSessions database table'
)

def handle(self, *args, **options):
User = get_user_model()
Expand All @@ -32,7 +34,9 @@ def handle(self, *args, **options):
for old_session in old_sessions:
if not UserSession.objects.filter(session_key=old_session.session_key).exists():
data = old_session.get_decoded()
user = User.objects.filter(pk=data['_auth_user_id']).first() if '_auth_user_id' in data else None
user = None
if '_auth_user_id' in data:
user = User.objects.filter(pk=data['_auth_user_id']).first()
UserSession.objects.create(
session_key=old_session.session_key,
session_data=old_session.session_data,
Expand Down

0 comments on commit b745418

Please sign in to comment.