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