1#!/usr/bin/env python
2# Copyright (c) 2012 Amazon.com, Inc. or its affiliates.  All Rights Reserved
3#
4# Permission is hereby granted, free of charge, to any person obtaining a
5# copy of this software and associated documentation files (the
6# "Software"), to deal in the Software without restriction, including
7# without limitation the rights to use, copy, modify, merge, publish, dis-
8# tribute, sublicense, and/or sell copies of the Software, and to permit
9# persons to whom the Software is furnished to do so, subject to the fol-
10# lowing conditions:
11#
12# The above copyright notice and this permission notice shall be included
13# in all copies or substantial portions of the Software.
14#
15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
16# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
17# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
18# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
19# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
21# IN THE SOFTWARE.
22#
23import time
24
25import boto
26from boto.compat import six
27from tests.compat import unittest
28from boto.ec2.networkinterface import NetworkInterfaceCollection
29from boto.ec2.networkinterface import NetworkInterfaceSpecification
30from boto.ec2.networkinterface import PrivateIPAddress
31
32
33class TestVPCConnection(unittest.TestCase):
34
35    def setUp(self):
36        # Registry of instances to be removed
37        self.instances = []
38        # Registry for cleaning up the vpc after all instances are terminated
39        # in the format [ ( func, (arg1, ... argn) ) ]
40        self.post_terminate_cleanups = []
41
42        self.api = boto.connect_vpc()
43        self.vpc = self.api.create_vpc('10.0.0.0/16')
44
45        # Need time for the VPC to be in place. :/
46        time.sleep(5)
47        self.subnet = self.api.create_subnet(self.vpc.id, '10.0.0.0/24')
48        # Register the subnet to be deleted after instance termination
49        self.post_terminate_cleanups.append((self.api.delete_subnet, (self.subnet.id,)))
50
51        # Need time for the subnet to be in place.
52        time.sleep(10)
53
54    def post_terminate_cleanup(self):
55        """Helper to run clean up tasks after instances are removed."""
56        for fn, args in self.post_terminate_cleanups:
57            fn(*args)
58            # Give things time to catch up each time
59            time.sleep(10)
60
61        # Now finally delete the vpc
62        if self.vpc:
63            self.api.delete_vpc(self.vpc.id)
64
65    def terminate_instances(self):
66        """Helper to remove all instances and kick off additional cleanup
67        once they are terminated.
68        """
69        for instance in self.instances:
70            self.terminate_instance(instance)
71        self.post_terminate_cleanup()
72
73    def terminate_instance(self, instance):
74        instance.terminate()
75        for i in six.moves.range(300):
76            instance.update()
77            if instance.state == 'terminated':
78                # Give it a litle more time to settle.
79                time.sleep(30)
80                return
81            else:
82                time.sleep(10)
83
84    def delete_elastic_ip(self, eip):
85        # Fetch a new copy of the eip so we're up to date
86        new_eip = self.api.get_all_addresses([eip.public_ip])[0]
87        if new_eip.association_id:
88            new_eip.disassociate()
89        new_eip.release()
90        time.sleep(10)
91
92    def test_multi_ip_create(self):
93        interface = NetworkInterfaceSpecification(
94            device_index=0, subnet_id=self.subnet.id,
95            private_ip_address='10.0.0.21',
96            description="This is a test interface using boto.",
97            delete_on_termination=True, private_ip_addresses=[
98                PrivateIPAddress(private_ip_address='10.0.0.22',
99                                 primary=False),
100                PrivateIPAddress(private_ip_address='10.0.0.23',
101                                 primary=False),
102                PrivateIPAddress(private_ip_address='10.0.0.24',
103                                 primary=False)])
104        interfaces = NetworkInterfaceCollection(interface)
105
106        reservation = self.api.run_instances(image_id='ami-a0cd60c9', instance_type='m1.small',
107                                             network_interfaces=interfaces)
108        # Give it a few seconds to start up.
109        time.sleep(10)
110        instance = reservation.instances[0]
111        self.addCleanup(self.terminate_instance, instance)
112        retrieved = self.api.get_all_reservations(instance_ids=[instance.id])
113        self.assertEqual(len(retrieved), 1)
114        retrieved_instances = retrieved[0].instances
115        self.assertEqual(len(retrieved_instances), 1)
116        retrieved_instance = retrieved_instances[0]
117
118        self.assertEqual(len(retrieved_instance.interfaces), 1)
119        interface = retrieved_instance.interfaces[0]
120
121        private_ip_addresses = interface.private_ip_addresses
122        self.assertEqual(len(private_ip_addresses), 4)
123        self.assertEqual(private_ip_addresses[0].private_ip_address,
124                         '10.0.0.21')
125        self.assertEqual(private_ip_addresses[0].primary, True)
126        self.assertEqual(private_ip_addresses[1].private_ip_address,
127                         '10.0.0.22')
128        self.assertEqual(private_ip_addresses[2].private_ip_address,
129                         '10.0.0.23')
130        self.assertEqual(private_ip_addresses[3].private_ip_address,
131                         '10.0.0.24')
132
133    def test_associate_public_ip(self):
134        # Supplying basically nothing ought to work.
135        interface = NetworkInterfaceSpecification(
136            associate_public_ip_address=True,
137            subnet_id=self.subnet.id,
138            # Just for testing.
139            delete_on_termination=True
140        )
141        interfaces = NetworkInterfaceCollection(interface)
142
143        reservation = self.api.run_instances(
144            image_id='ami-a0cd60c9',
145            instance_type='m1.small',
146            network_interfaces=interfaces
147        )
148        instance = reservation.instances[0]
149        self.instances.append(instance)
150        self.addCleanup(self.terminate_instances)
151
152        # Give it a **LONG** time to start up.
153        # Because the public IP won't be there right away.
154        time.sleep(60)
155
156        retrieved = self.api.get_all_reservations(
157            instance_ids=[
158                instance.id
159            ]
160        )
161        self.assertEqual(len(retrieved), 1)
162        retrieved_instances = retrieved[0].instances
163        self.assertEqual(len(retrieved_instances), 1)
164        retrieved_instance = retrieved_instances[0]
165
166        self.assertEqual(len(retrieved_instance.interfaces), 1)
167        interface = retrieved_instance.interfaces[0]
168
169        # There ought to be a public IP there.
170        # We can't reason about the IP itself, so just make sure it vaguely
171        # resembles an IP (& isn't empty/``None``)...
172        self.assertTrue(interface.publicIp.count('.') >= 3)
173
174    def test_associate_elastic_ip(self):
175        interface = NetworkInterfaceSpecification(
176            associate_public_ip_address=False,
177            subnet_id=self.subnet.id,
178            # Just for testing.
179            delete_on_termination=True
180        )
181        interfaces = NetworkInterfaceCollection(interface)
182
183        reservation = self.api.run_instances(
184            image_id='ami-a0cd60c9',
185            instance_type='m1.small',
186            network_interfaces=interfaces
187        )
188        instance = reservation.instances[0]
189        # Register instance to be removed
190        self.instances.append(instance)
191        # Add terminate instances helper as cleanup command
192        self.addCleanup(self.terminate_instances)
193
194        # Create an internet gateway so we can attach an eip
195        igw = self.api.create_internet_gateway()
196        # Wait on gateway before attaching
197        time.sleep(5)
198        # Attach and register clean up tasks
199        self.api.attach_internet_gateway(igw.id, self.vpc.id)
200        self.post_terminate_cleanups.append((self.api.detach_internet_gateway, (igw.id, self.vpc.id)))
201        self.post_terminate_cleanups.append((self.api.delete_internet_gateway, (igw.id,)))
202
203        # Allocate an elastic ip to this vpc
204        eip = self.api.allocate_address('vpc')
205        self.post_terminate_cleanups.append((self.delete_elastic_ip, (eip,)))
206
207        # Wait on instance and eip then try to associate directly to instance
208        time.sleep(60)
209        eip.associate(instance.id)
210
211
212if __name__ == '__main__':
213    unittest.main()
214