| 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 |