1from webapp2_extras import auth
2from webapp2_extras.appengine.auth import models
3
4from google.appengine.ext.ndb import model
5
6import test_base
7
8
9class UniqueConstraintViolation(Exception):
10    pass
11
12
13class User(model.Model):
14    username = model.StringProperty(required=True)
15    auth_id = model.StringProperty()
16    email = model.StringProperty()
17
18
19class TestAuthModels(test_base.BaseTestCase):
20
21    def setUp(self):
22        super(TestAuthModels, self).setUp()
23        self.register_model('User', models.User)
24        self.register_model('UserToken', models.UserToken)
25        self.register_model('Unique', models.Unique)
26
27    def test_get(self):
28        m = models.User
29        success, user = m.create_user(auth_id='auth_id_1', password_raw='foo')
30        self.assertEqual(success, True)
31        self.assertTrue(user is not None)
32        self.assertTrue(user.password is not None)
33
34        # user.key.id() is required to retrieve the auth token
35        user_id = user.key.id()
36
37        token = m.create_auth_token(user_id)
38
39        self.assertEqual(m.get_by_auth_id('auth_id_1'), user)
40        self.assertEqual(m.get_by_auth_id('auth_id_2'), None)
41
42        u, ts = m.get_by_auth_token(user_id, token)
43        self.assertEqual(u, user)
44        u, ts = m.get_by_auth_token('fake_user_id', token)
45        self.assertEqual(u, None)
46
47        u = m.get_by_auth_password('auth_id_1', 'foo')
48        self.assertEqual(u, user)
49        self.assertRaises(auth.InvalidPasswordError,
50                          m.get_by_auth_password, 'auth_id_1', 'bar')
51        self.assertRaises(auth.InvalidAuthIdError,
52                          m.get_by_auth_password, 'auth_id_2', 'foo')
53
54    def test_create_user(self):
55        m = models.User
56        success, info = m.create_user(auth_id='auth_id_1', password_raw='foo')
57        self.assertEqual(success, True)
58        self.assertTrue(info is not None)
59        self.assertTrue(info.password is not None)
60
61        success, info = m.create_user(auth_id='auth_id_1')
62        self.assertEqual(success, False)
63        self.assertEqual(info, ['auth_id'])
64
65        # 3 extras and unique properties; plus 1 extra and not unique.
66        extras = ['foo', 'bar', 'baz']
67        values = dict((v, v + '_value') for v in extras)
68        values['ding'] = 'ding_value'
69        success, info = m.create_user(auth_id='auth_id_2',
70                                      unique_properties=extras, **values)
71        self.assertEqual(success, True)
72        self.assertTrue(info is not None)
73        for prop in extras:
74            self.assertEqual(getattr(info, prop), prop + '_value')
75        self.assertEqual(info.ding, 'ding_value')
76
77        # Let's do it again.
78        success, info = m.create_user(auth_id='auth_id_3',
79                                      unique_properties=extras, **values)
80        self.assertEqual(success, False)
81        self.assertEqual(info, extras)
82
83    def test_add_auth_ids(self):
84        m = models.User
85        success, new_user = m.create_user(auth_id='auth_id_1', password_raw='foo')
86        self.assertEqual(success, True)
87        self.assertTrue(new_user is not None)
88        self.assertTrue(new_user.password is not None)
89
90        success, new_user_2 = m.create_user(auth_id='auth_id_2', password_raw='foo')
91        self.assertEqual(success, True)
92        self.assertTrue(new_user is not None)
93        self.assertTrue(new_user.password is not None)
94
95        success, info = new_user.add_auth_id('auth_id_3')
96        self.assertEqual(success, True)
97        self.assertEqual(info.auth_ids, ['auth_id_1', 'auth_id_3'])
98
99        success, info = new_user.add_auth_id('auth_id_2')
100        self.assertEqual(success, False)
101        self.assertEqual(info, ['auth_id'])
102
103    def test_token(self):
104        m = models.UserToken
105
106        auth_id = 'foo'
107        subject = 'bar'
108        token_1 = m.create(auth_id, subject, token=None)
109        token = token_1.token
110
111        token_2 = m.get(user=auth_id, subject=subject, token=token)
112        self.assertEqual(token_2, token_1)
113
114        token_3 = m.get(subject=subject, token=token)
115        self.assertEqual(token_3, token_1)
116
117        m.get_key(auth_id, subject, token).delete()
118
119        token_2 = m.get(user=auth_id, subject=subject, token=token)
120        self.assertEqual(token_2, None)
121
122        token_3 = m.get(subject=subject, token=token)
123        self.assertEqual(token_3, None)
124
125    def test_user_token(self):
126        m = models.User
127        auth_id = 'foo'
128
129        token = m.create_auth_token(auth_id)
130        self.assertTrue(m.validate_auth_token(auth_id, token))
131        m.delete_auth_token(auth_id, token)
132        self.assertFalse(m.validate_auth_token(auth_id, token))
133
134        token = m.create_signup_token(auth_id)
135        self.assertTrue(m.validate_signup_token(auth_id, token))
136        m.delete_signup_token(auth_id, token)
137        self.assertFalse(m.validate_signup_token(auth_id, token))
138
139
140class TestUniqueModel(test_base.BaseTestCase):
141    def setUp(self):
142        super(TestUniqueModel, self).setUp()
143        self.register_model('Unique', models.Unique)
144
145    def test_single(self):
146        def create_user(username):
147            # Assemble the unique scope/value combinations.
148            unique_username = 'User.username:%s' % username
149
150            # Create the unique username, auth_id and email.
151            success = models.Unique.create(unique_username)
152
153            if success:
154                user = User(username=username)
155                user.put()
156                return user
157            else:
158                raise UniqueConstraintViolation('Username %s already '
159                    'exists' % username)
160
161        user = create_user('username_1')
162        self.assertRaises(UniqueConstraintViolation, create_user, 'username_1')
163
164        user = create_user('username_2')
165        self.assertRaises(UniqueConstraintViolation, create_user, 'username_2')
166
167    def test_multi(self):
168        def create_user(username, auth_id, email):
169            # Assemble the unique scope/value combinations.
170            unique_username = 'User.username:%s' % username
171            unique_auth_id = 'User.auth_id:%s' % auth_id
172            unique_email = 'User.email:%s' % email
173
174            # Create the unique username, auth_id and email.
175            uniques = [unique_username, unique_auth_id, unique_email]
176            success, existing = models.Unique.create_multi(uniques)
177
178            if success:
179                user = User(username=username, auth_id=auth_id, email=email)
180                user.put()
181                return user
182            else:
183                if unique_username in existing:
184                    raise UniqueConstraintViolation('Username %s already '
185                        'exists' % username)
186                if unique_auth_id in existing:
187                    raise UniqueConstraintViolation('Auth id %s already '
188                        'exists' % auth_id)
189                if unique_email in existing:
190                    raise UniqueConstraintViolation('Email %s already '
191                        'exists' % email)
192
193        user = create_user('username_1', 'auth_id_1', 'email_1')
194        self.assertRaises(UniqueConstraintViolation, create_user, 'username_1', 'auth_id_2', 'email_2')
195        self.assertRaises(UniqueConstraintViolation, create_user, 'username_2', 'auth_id_1', 'email_2')
196        self.assertRaises(UniqueConstraintViolation, create_user, 'username_2', 'auth_id_2', 'email_1')
197
198        user = create_user('username_2', 'auth_id_2', 'email_2')
199        self.assertRaises(UniqueConstraintViolation, create_user, 'username_2', 'auth_id_1', 'email_1')
200        self.assertRaises(UniqueConstraintViolation, create_user, 'username_1', 'auth_id_2', 'email_1')
201        self.assertRaises(UniqueConstraintViolation, create_user, 'username_1', 'auth_id_1', 'email_2')
202
203    def test_delete_multi(self):
204        rv = models.Unique.create_multi(('foo', 'bar', 'baz'))
205        self.assertEqual(rv, (True, []))
206        rv = models.Unique.create_multi(('foo', 'bar', 'baz'))
207        self.assertEqual(rv, (False, ['foo', 'bar', 'baz']))
208
209        models.Unique.delete_multi(('foo', 'bar', 'baz'))
210
211        rv = models.Unique.create_multi(('foo', 'bar', 'baz'))
212        self.assertEqual(rv, (True, []))
213
214
215if __name__ == '__main__':
216    test_base.main()
217