1#!/usr/bin/env python
2#
3# Copyright 2016 - The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""Tests for acloud.public.device_driver."""
18
19import datetime
20import uuid
21
22import dateutil.parser
23import mock
24
25import unittest
26from acloud.internal.lib import auth
27from acloud.internal.lib import android_build_client
28from acloud.internal.lib import android_compute_client
29from acloud.internal.lib import driver_test_lib
30from acloud.internal.lib import gstorage_client
31from acloud.public import device_driver
32
33
34class DeviceDriverTest(driver_test_lib.BaseDriverTest):
35    """Test device_driver."""
36
37    def setUp(self):
38        """Set up the test."""
39        super(DeviceDriverTest, self).setUp()
40        self.build_client = mock.MagicMock()
41        self.Patch(android_build_client, "AndroidBuildClient",
42            return_value=self.build_client)
43        self.storage_client = mock.MagicMock()
44        self.Patch(
45            gstorage_client, "StorageClient", return_value=self.storage_client)
46        self.compute_client = mock.MagicMock()
47        self.Patch(
48            android_compute_client,
49            "AndroidComputeClient",
50            return_value=self.compute_client)
51        self.Patch(auth, "CreateCredentials", return_value=mock.MagicMock())
52
53    def _CreateCfg(self):
54        """A helper method that creates a mock configuration object."""
55        cfg = mock.MagicMock()
56        cfg.service_account_name = "fake@service.com"
57        cfg.service_account_private_key_path = "/fake/path/to/key"
58        cfg.zone = "fake_zone"
59        cfg.disk_image_name = "fake_image.tar.gz"
60        cfg.disk_image_mime_type = "fake/type"
61        cfg.storage_bucket_name = "fake_bucket"
62        cfg.extra_data_disk_size_gb = 4
63        cfg.precreated_data_image_map = {
64            4: "extradisk-image-4gb",
65            10: "extradisk-image-10gb"
66        }
67        cfg.ssh_private_key_path = ""
68        cfg.ssh_public_key_path = ""
69
70        return cfg
71
72    def testCreateAndroidVirtualDevices(self):
73        """Test CreateAndroidVirtualDevices."""
74        cfg = self._CreateCfg()
75        fake_gs_url = "fake_gs_url"
76        fake_ip = "140.1.1.1"
77        fake_instance = "fake-instance"
78        fake_image = "fake-image"
79        fake_build_target = "fake_target"
80        fake_build_id = "12345"
81
82        # Mock uuid
83        fake_uuid = mock.MagicMock(hex="1234")
84        self.Patch(uuid, "uuid4", return_value=fake_uuid)
85        fake_gs_object = fake_uuid.hex + "-" + cfg.disk_image_name
86        self.storage_client.GetUrl.return_value = fake_gs_url
87
88        # Mock compute client methods
89        disk_name = "extradisk-image-4gb"
90        self.compute_client.GetInstanceIP.return_value = fake_ip
91        self.compute_client.GenerateImageName.return_value = fake_image
92        self.compute_client.GenerateInstanceName.return_value = fake_instance
93        self.compute_client.GetDataDiskName.return_value = disk_name
94
95        # Verify
96        r = device_driver.CreateAndroidVirtualDevices(
97        cfg, fake_build_target, fake_build_id)
98        self.build_client.CopyTo.assert_called_with(
99        fake_build_target, fake_build_id, artifact_name=cfg.disk_image_name,
100        destination_bucket=cfg.storage_bucket_name,
101        destination_path=fake_gs_object)
102        self.compute_client.CreateImage.assert_called_with(
103        image_name=fake_image, source_uri=fake_gs_url)
104        self.compute_client.CreateInstance.assert_called_with(
105        fake_instance, fake_image, disk_name)
106        self.compute_client.DeleteImage.assert_called_with(fake_image)
107        self.storage_client.Delete(cfg.storage_bucket_name, fake_gs_object)
108
109        self.assertEquals(
110            r.data,
111            {
112                "devices": [
113                    {
114                        "instance_name": fake_instance,
115                        "ip": fake_ip,
116                    },
117                ],
118            }
119        )
120        self.assertEquals(r.command, "create")
121        self.assertEquals(r.status, "SUCCESS")
122
123
124    def testDeleteAndroidVirtualDevices(self):
125        """Test DeleteAndroidVirtualDevices."""
126        instance_names = ["fake-instance-1", "fake-instance-2"]
127        self.compute_client.DeleteInstances.return_value = (instance_names, [],
128                                                            [])
129        cfg = self._CreateCfg()
130        r = device_driver.DeleteAndroidVirtualDevices(cfg, instance_names)
131        self.compute_client.DeleteInstances.assert_called_once_with(
132            instance_names, cfg.zone)
133        self.assertEquals(r.data, {
134            "deleted": [
135                {
136                    "name": instance_names[0],
137                    "type": "instance",
138                },
139                {
140                    "name": instance_names[1],
141                    "type": "instance",
142                },
143            ],
144        })
145        self.assertEquals(r.command, "delete")
146        self.assertEquals(r.status, "SUCCESS")
147
148    def testCleanup(self):
149        expiration_mins = 30
150        before_deadline = "2015-10-29T12:00:30.018-07:00"
151        after_deadline = "2015-10-29T12:45:30.018-07:00"
152        now = "2015-10-29T13:00:30.018-07:00"
153        self.Patch(device_driver, "datetime")
154        device_driver.datetime.datetime.now.return_value = dateutil.parser.parse(
155            now)
156        device_driver.datetime.timedelta.return_value = datetime.timedelta(
157            minutes=expiration_mins)
158        fake_instances = [
159            {
160                "name": "fake_instance_1",
161                "creationTimestamp": before_deadline,
162            }, {
163                "name": "fake_instance_2",
164                "creationTimestamp": after_deadline,
165            }
166        ]
167        fake_images = [
168            {
169                "name": "extradisk-image-4gb",
170                "creationTimestamp": before_deadline,
171            }, {
172                "name": "fake_image_1",
173                "creationTimestamp": before_deadline,
174            }, {
175                "name": "fake_image_2",
176                "creationTimestamp": after_deadline,
177            }
178        ]
179        fake_disks = [
180            {
181                "name": "fake_disk_1",
182                "creationTimestamp": before_deadline,
183            }, {
184                "name": "fake_disk_2",
185                "creationTimestamp": before_deadline,
186                "users": ["some-instance-using-the-disk"]
187            }, {
188                "name": "fake_disk_3",
189                "creationTimestamp": after_deadline,
190            }
191        ]
192        fake_objects = [
193            {
194                "name": "fake_object_1",
195                "timeCreated": before_deadline,
196            }, {
197                "name": "fake_object_2",
198                "timeCreated": after_deadline,
199            }
200        ]
201        self.compute_client.ListInstances.return_value = fake_instances
202        self.compute_client.ListImages.return_value = fake_images
203        self.compute_client.ListDisks.return_value = fake_disks
204        self.storage_client.List.return_value = fake_objects
205        self.compute_client.DeleteInstances.return_value = (
206            ["fake_instance_1"], [], [])
207        self.compute_client.DeleteImages.return_value = (["fake_image_1"], [],
208                                                         [])
209        self.compute_client.DeleteDisks.return_value = (["fake_disk_1"], [],
210                                                        [])
211        self.storage_client.DeleteFiles.return_value = (["fake_object_1"], [],
212                                                        [])
213        cfg = self._CreateCfg()
214        r = device_driver.Cleanup(cfg, expiration_mins)
215        self.assertEqual(r.errors, [])
216        expected_report_data = {
217            "deleted": [
218                {"name": "fake_instance_1",
219                 "type": "instance"},
220                {"name": "fake_image_1",
221                 "type": "image"},
222                {"name": "fake_disk_1",
223                 "type": "disk"},
224                {"name": "fake_object_1",
225                 "type": "cached_build_artifact"},
226            ]
227        }
228        self.assertEqual(r.data, expected_report_data)
229
230        self.compute_client.ListInstances.assert_called_once_with(
231            zone=cfg.zone)
232        self.compute_client.DeleteInstances.assert_called_once_with(
233            instances=["fake_instance_1"], zone=cfg.zone)
234
235        self.compute_client.ListImages.assert_called_once_with()
236        self.compute_client.DeleteImages.assert_called_once_with(
237            image_names=["fake_image_1"])
238
239        self.compute_client.ListDisks.assert_called_once_with(zone=cfg.zone)
240        self.compute_client.DeleteDisks.assert_called_once_with(
241            disk_names=["fake_disk_1"], zone=cfg.zone)
242
243        self.storage_client.List.assert_called_once_with(
244            bucket_name=cfg.storage_bucket_name)
245        self.storage_client.DeleteFiles.assert_called_once_with(
246            bucket_name=cfg.storage_bucket_name,
247            object_names=["fake_object_1"])
248
249
250if __name__ == "__main__":
251    unittest.main()
252