1#!/usr/bin/env python
2#
3# Copyright (C) 2018 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
18import os
19import unittest
20
21try:
22    from unittest import mock
23except ImportError:
24    # TODO: Remove when we stop supporting Python 2
25    import mock
26
27from google.auth import exceptions as auth_exceptions
28
29from vts.utils.python.gcs import gcs_api_utils
30
31
32def simple_ListFilesWithPrefix(dir_path):
33    return [
34        '%s/file1' % dir_path,
35        '%s/file2' % dir_path,
36        '%s/file3' % dir_path,
37        '%s/file4' % dir_path
38    ]
39
40
41def simple_DownloadFile(src_file_path, dest_file_path):
42    return None
43
44
45def simple_UploadFile(src_file_path, dest_file_path):
46    return None
47
48
49def simple_PrefixExists(dir_path):
50    if dir_path is 'valid_source_dir':
51        return True
52    else:
53        return False
54
55
56def simple_os_path_exists(path):
57    return True
58
59
60def simple__init__(key_path, bucket_name):
61    return None
62
63
64def simple_PrepareDownloadDestination(src_dir, dest_dir):
65    return os.path.join(dest_dir, os.path.basename(src_dir))
66
67
68class GcsApiUtilsTest(unittest.TestCase):
69    """Unit tests for gcs_utils module."""
70
71    def SetUp(self):
72        """Setup tasks."""
73        self.category = "category_default"
74        self.name = "name_default"
75
76    @mock.patch(
77        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.ListFilesWithPrefix',
78        side_effect=simple_ListFilesWithPrefix)
79    @mock.patch(
80        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.__init__',
81        side_effect=simple__init__)
82    def testCountFiles(self, simple__init__, simple_ListFilesWithPrefix):
83        """Tests the CountFiles function."""
84        _gcs_api_utils = gcs_api_utils.GcsApiUtils('key/path', 'vts-fuzz')
85        length = _gcs_api_utils.CountFiles('corpus/ILight/ILight_corpus_seed')
86        simple_ListFilesWithPrefix.assert_called()
87        self.assertEqual(length, 4)
88
89    @mock.patch(
90        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.ListFilesWithPrefix',
91        side_effect=simple_ListFilesWithPrefix)
92    @mock.patch(
93        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.__init__',
94        side_effect=simple__init__)
95    def testPrefixExists(self, simple__init__, simple_ListFilesWithPrefix):
96        """Tests the PrefixExists function."""
97        _gcs_api_utils = gcs_api_utils.GcsApiUtils('key/path', 'vts-fuzz')
98        dir_exist = _gcs_api_utils.PrefixExists(
99            'corpus/ILight/ILight_corpus_seed')
100        simple_ListFilesWithPrefix.assert_called()
101        self.assertEqual(dir_exist, True)
102
103    @mock.patch('os.path.exists', side_effect=simple_os_path_exists)
104    @mock.patch(
105        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.__init__',
106        side_effect=simple__init__)
107    def testPrepareDownloadDestination(self, simple__init__,
108                                       simple_os_path_exists):
109        """Tests the PrepareDownloadDestination function."""
110        _gcs_api_utils = gcs_api_utils.GcsApiUtils('key/path', 'vts-fuzz')
111        local_dest_folder = _gcs_api_utils.PrepareDownloadDestination(
112            'corpus/ILight/ILight_corpus_seed', 'tmp/tmp4772')
113        self.assertEqual(local_dest_folder, 'tmp/tmp4772/ILight_corpus_seed')
114
115    @mock.patch(
116        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.DownloadFile',
117        side_effect=simple_DownloadFile)
118    @mock.patch(
119        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.PrefixExists',
120        side_effect=simple_PrefixExists)
121    @mock.patch(
122        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.ListFilesWithPrefix',
123        side_effect=simple_ListFilesWithPrefix)
124    @mock.patch(
125        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.PrepareDownloadDestination',
126        side_effect=simple_PrepareDownloadDestination)
127    @mock.patch(
128        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.__init__',
129        side_effect=simple__init__)
130    def testDownloadDir(self, simple__init__,
131                        simple_PrepareDownloadDestination,
132                        simple_ListFilesWithPrefix, simple_PrefixExists,
133                        simple_DownloadFile):
134        """Tests the DownloadDir function"""
135        _gcs_api_utils = gcs_api_utils.GcsApiUtils('key/path', 'vts-fuzz')
136        _gcs_api_utils.DownloadDir('valid_source_dir',
137                                   'local_destination/dest')
138        num_DownloadFile_called = simple_DownloadFile.call_count
139        self.assertEqual(num_DownloadFile_called, 4)
140        local_dest_folder = simple_PrepareDownloadDestination.return_value
141
142    @mock.patch(
143        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.UploadFile',
144        side_effect=simple_UploadFile)
145    @mock.patch('os.path.exists', side_effect=simple_os_path_exists)
146    @mock.patch('os.listdir', side_effect=simple_ListFilesWithPrefix)
147    @mock.patch(
148        'vts.utils.python.gcs.gcs_api_utils.GcsApiUtils.__init__',
149        side_effect=simple__init__)
150    def testUploadDir(self, simple__init__, simple_ListFilesWithPrefix,
151                      simple_os_path_exists, simple_UploadFile):
152        _gcs_api_utils = gcs_api_utils.GcsApiUtils('key/path', 'vts-fuzz')
153        _gcs_api_utils.UploadDir('valid_source_dir', 'GCS_destination/dest')
154        num_UploadFile_called = simple_UploadFile.call_count
155        self.assertEqual(num_UploadFile_called, 4)
156
157    @mock.patch(
158        'vts.utils.python.gcs.gcs_api_utils.google.auth.default',
159        side_effect=auth_exceptions.DefaultCredentialsError('unit test'))
160    def testCredentialsError(self, mock_default):
161        """Tests authentication failure in __init__."""
162        _gcs_api_utils = gcs_api_utils.GcsApiUtils('key/path', 'vts-fuzz')
163        self.assertFalse(_gcs_api_utils.Enabled)
164        mock_default.assert_called()
165
166
167if __name__ == "__main__":
168    unittest.main()
169