OLD | NEW |
| (Empty) |
1 #!/usr/bin/python | |
2 # Copyright 2016 The LUCI Authors. All rights reserved. | |
3 # Use of this source code is governed under the Apache License, Version 2.0 | |
4 # that can be found in the LICENSE file. | |
5 | |
6 """Unit tests for pubsub.py.""" | |
7 | |
8 import base64 | |
9 import unittest | |
10 | |
11 import test_env | |
12 test_env.setup_test_env() | |
13 | |
14 from google.appengine.ext import ndb | |
15 | |
16 from components import datastore_utils | |
17 from test_support import test_case | |
18 | |
19 import instances | |
20 import models | |
21 import pubsub | |
22 | |
23 | |
24 class ProcessTest(test_case.TestCase): | |
25 """Tests for pubsub.process.""" | |
26 | |
27 def setUp(self): | |
28 super(ProcessTest, self).setUp() | |
29 | |
30 def ack_async(*args, **kwargs): | |
31 return ndb.Future() | |
32 self.mock(pubsub.pubsub, 'ack_async', ack_async) | |
33 | |
34 def test_leased(self): | |
35 """Ensures nothing happens when an instance is leased.""" | |
36 key = models.Instance( | |
37 key=instances.get_instance_key( | |
38 'base-name', | |
39 'revision', | |
40 'zone', | |
41 'instance-name', | |
42 ), | |
43 ).put() | |
44 message = { | |
45 'ackId': 'id', | |
46 'message': { | |
47 'attributes': { | |
48 'key': key.urlsafe(), | |
49 }, | |
50 'data': base64.b64encode('LEASED'), | |
51 }, | |
52 } | |
53 | |
54 pubsub.process(message).wait() | |
55 self.failUnless(key.get()) | |
56 | |
57 def test_reclaimed(self): | |
58 """Ensures reclaimed instance is marked for deletion.""" | |
59 key = models.Instance( | |
60 key=instances.get_instance_key( | |
61 'base-name', | |
62 'revision', | |
63 'zone', | |
64 'instance-name', | |
65 ), | |
66 ).put() | |
67 message = { | |
68 'ackId': 'id', | |
69 'message': { | |
70 'attributes': { | |
71 'key': key.urlsafe(), | |
72 }, | |
73 'data': base64.b64encode('RECLAIMED'), | |
74 }, | |
75 } | |
76 | |
77 pubsub.process(message).wait() | |
78 self.failUnless(key.get().pending_deletion) | |
79 | |
80 def test_subscribed(self): | |
81 """Ensures subscribed instance has pending metadata update.""" | |
82 key = models.Instance( | |
83 key=instances.get_instance_key( | |
84 'base-name', | |
85 'revision', | |
86 'zone', | |
87 'instance-name', | |
88 ), | |
89 ).put() | |
90 models.InstanceGroupManager( | |
91 key=key.parent(), | |
92 ).put() | |
93 models.InstanceTemplateRevision( | |
94 key=key.parent().parent(), | |
95 service_accounts=[ | |
96 models.ServiceAccount( | |
97 name='name', | |
98 scopes=[ | |
99 'scope', | |
100 ], | |
101 ), | |
102 ], | |
103 ).put() | |
104 message = { | |
105 'ackId': 'id', | |
106 'message': { | |
107 'attributes': { | |
108 'key': key.urlsafe(), | |
109 'subscription': 'subscription', | |
110 'subscription_project': 'subscription-project', | |
111 }, | |
112 'data': base64.b64encode('SUBSCRIBED'), | |
113 }, | |
114 } | |
115 expected_pending_metadata_updates = [ | |
116 models.MetadataUpdate( | |
117 metadata={ | |
118 'pubsub_service_account': 'name', | |
119 'pubsub_subscription': 'subscription', | |
120 'pubsub_subscription_project': 'subscription-project', | |
121 }, | |
122 ) | |
123 ] | |
124 | |
125 pubsub.process(message).wait() | |
126 self.assertEqual( | |
127 key.get().pending_metadata_updates, expected_pending_metadata_updates) | |
128 | |
129 def test_unexpected_key(self): | |
130 """Ensures nothing happens when key has unexpected kind.""" | |
131 key = models.Instance( | |
132 key=instances.get_instance_key( | |
133 'base-name', | |
134 'revision', | |
135 'zone', | |
136 'instance-name', | |
137 ), | |
138 ).put() | |
139 models.InstanceGroupManager( | |
140 key=key.parent(), | |
141 ).put() | |
142 models.InstanceTemplateRevision( | |
143 key=key.parent().parent(), | |
144 service_accounts=[ | |
145 models.ServiceAccount( | |
146 name='name', | |
147 scopes=[ | |
148 'scope', | |
149 ], | |
150 ), | |
151 ], | |
152 ).put() | |
153 message = { | |
154 'ackId': 'id', | |
155 'message': { | |
156 'attributes': { | |
157 'key': key.parent().urlsafe(), | |
158 'subscription': 'subscription', | |
159 'subscription_project': 'subscription-project', | |
160 }, | |
161 'data': base64.b64encode('SUBSCRIBED'), | |
162 }, | |
163 } | |
164 | |
165 pubsub.process(message).wait() | |
166 self.failIf(key.get().pending_metadata_updates) | |
167 | |
168 | |
169 if __name__ == '__main__': | |
170 unittest.main() | |
OLD | NEW |