1# Copyright 2015 The Chromium Authors. All rights reserved. 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5import unittest 6 7import mock 8 9from google.appengine.ext import ndb 10 11from dashboard import testing_common 12from dashboard.models import internal_only_model 13 14 15class InternalOnlyModelExample(internal_only_model.InternalOnlyModel): 16 """Example Model which has an internal_only property.""" 17 18 internal_only = ndb.BooleanProperty(default=False) 19 20 21class CreateHookModelExample(internal_only_model.CreateHookInternalOnlyModel): 22 """Example Model which has a CreateHook method.""" 23 24 internal_only = ndb.BooleanProperty(default=False) 25 26 def CreateCallback(self): 27 pass 28 29 30class InternalOnlyModelTest(testing_common.TestCase): 31 32 def setUp(self): 33 super(InternalOnlyModelTest, self).setUp() 34 testing_common.SetIsInternalUser('x@google.com', True) 35 testing_common.SetIsInternalUser('x@foo.com', False) 36 37 def testInternalOnlyModel_InternalUser_EntityFetched(self): 38 key = InternalOnlyModelExample(internal_only=True).put() 39 self.SetCurrentUser('x@google.com') 40 self.assertEqual(key, key.get().key) 41 42 def testInternalOnlyModel_ExternalUserInternalEntity_AssertionRaised(self): 43 key = InternalOnlyModelExample(internal_only=True).put() 44 self.SetCurrentUser('x@foo.com') 45 with self.assertRaises(AssertionError): 46 key.get() 47 48 def testInternalOnlyModel_ExternalUserExternalEntity_EntityFetched(self): 49 key = InternalOnlyModelExample(internal_only=False).put() 50 self.SetCurrentUser('x@foo.com') 51 self.assertEqual(key, key.get().key) 52 53 @mock.patch.object(CreateHookModelExample, 'CreateCallback') 54 def testCreateCallBack_CalledOncePerEntity(self, mock_create_callback): 55 one = CreateHookModelExample() 56 two = CreateHookModelExample() 57 self.assertEqual(0, mock_create_callback.call_count) 58 one_key = one.put() 59 one.put() 60 self.assertEqual(1, mock_create_callback.call_count) 61 two_key = two.put() 62 self.assertEqual(2, mock_create_callback.call_count) 63 one_key.get().put() 64 two_key.get().put() 65 self.assertEqual(2, mock_create_callback.call_count) 66 67 def testCreateHookEntity_InternalUserInternalEntity_EntityFetched(self): 68 key = CreateHookModelExample(internal_only=True).put() 69 self.SetCurrentUser('x@google.com') 70 self.assertEqual(key, key.get().key) 71 72 def testCreateHookEntity_ExternalUserInternalEntity_AssertionRaised(self): 73 key = CreateHookModelExample(internal_only=True).put() 74 self.SetCurrentUser('x@foo.com') 75 with self.assertRaises(AssertionError): 76 key.get() 77 78 def testCreateHookEntity_ExternalUserExternalEntity_EntityFetched(self): 79 key = CreateHookModelExample(internal_only=False).put() 80 self.SetCurrentUser('x@foo.com') 81 self.assertEqual(key, key.get().key) 82 83 84if __name__ == '__main__': 85 unittest.main() 86