1#!/usr/bin/env python3 2# 3# Copyright 2021, 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"""Unittests for atest_gcp_utils.""" 18 19import os 20import tempfile 21import unittest 22 23from pathlib import Path 24from unittest import mock 25 26import constants 27 28from logstorage import atest_gcp_utils 29 30class AtestGcpUtilsUnittests(unittest.TestCase): 31 """Unit tests for atest_gcp_utils.py""" 32 33 @mock.patch.object(atest_gcp_utils, '_prepare_data') 34 @mock.patch.object(atest_gcp_utils, 'fetch_credential') 35 def test_do_upload_flow(self, mock_request, mock_prepare): 36 """test do_upload_flow method.""" 37 fake_extra_args = {} 38 fake_creds = mock.Mock() 39 fake_creds.token_response = {'access_token': 'fake_token'} 40 mock_request.return_value = fake_creds 41 fake_inv = {'invocationId': 'inv_id'} 42 fake_workunit = {'id': 'workunit_id'} 43 fake_local_build_id = 'L1234567' 44 fake_build_target = 'build_target' 45 mock_prepare.return_value = (fake_inv, fake_workunit, 46 fake_local_build_id, fake_build_target) 47 constants.TOKEN_FILE_PATH = tempfile.NamedTemporaryFile().name 48 creds, inv = atest_gcp_utils.do_upload_flow(fake_extra_args) 49 self.assertEqual(fake_creds, creds) 50 self.assertEqual(fake_inv, inv) 51 self.assertEqual(fake_extra_args[constants.INVOCATION_ID], 52 fake_inv['invocationId']) 53 self.assertEqual(fake_extra_args[constants.WORKUNIT_ID], 54 fake_workunit['id']) 55 self.assertEqual(fake_extra_args[constants.LOCAL_BUILD_ID], 56 fake_local_build_id) 57 self.assertEqual(fake_extra_args[constants.BUILD_TARGET], 58 fake_build_target) 59 60 mock_request.return_value = None 61 creds, inv = atest_gcp_utils.do_upload_flow(fake_extra_args) 62 self.assertEqual(None, creds) 63 self.assertEqual(None, inv) 64 65 @mock.patch.object(atest_gcp_utils.GCPHelper, 66 'get_credential_with_auth_flow') 67 def test_fetch_credential_with_request_option( 68 self, mock_get_credential_with_auth_flow): 69 """test fetch_credential method.""" 70 constants.CREDENTIAL_FILE_NAME = 'cred_file' 71 constants.GCP_ACCESS_TOKEN = 'access_token' 72 tmp_folder = tempfile.mkdtemp() 73 not_upload_file = os.path.join(tmp_folder, 74 constants.DO_NOT_UPLOAD) 75 76 atest_gcp_utils.fetch_credential(tmp_folder, 77 {constants.REQUEST_UPLOAD_RESULT:True}) 78 self.assertEqual(1, mock_get_credential_with_auth_flow.call_count) 79 self.assertFalse(os.path.exists(not_upload_file)) 80 81 atest_gcp_utils.fetch_credential(tmp_folder, 82 {constants.REQUEST_UPLOAD_RESULT:True}) 83 self.assertEqual(2, mock_get_credential_with_auth_flow.call_count) 84 self.assertFalse(os.path.exists(not_upload_file)) 85 86 @mock.patch.object(atest_gcp_utils.GCPHelper, 87 'get_credential_with_auth_flow') 88 def test_fetch_credential_with_disable_option( 89 self, mock_get_credential_with_auth_flow): 90 """test fetch_credential method.""" 91 constants.CREDENTIAL_FILE_NAME = 'cred_file' 92 constants.GCP_ACCESS_TOKEN = 'access_token' 93 tmp_folder = tempfile.mkdtemp() 94 not_upload_file = os.path.join(tmp_folder, 95 constants.DO_NOT_UPLOAD) 96 97 atest_gcp_utils.fetch_credential(tmp_folder, 98 {constants.DISABLE_UPLOAD_RESULT:True}) 99 self.assertTrue(os.path.exists(not_upload_file)) 100 self.assertEqual(0, mock_get_credential_with_auth_flow.call_count) 101 102 atest_gcp_utils.fetch_credential(tmp_folder, 103 {constants.DISABLE_UPLOAD_RESULT:True}) 104 self.assertEqual(0, mock_get_credential_with_auth_flow.call_count) 105 106 @mock.patch.object(atest_gcp_utils.GCPHelper, 107 'get_credential_with_auth_flow') 108 def test_fetch_credential_no_upload_option( 109 self, mock_get_credential_with_auth_flow): 110 """test fetch_credential method.""" 111 constants.CREDENTIAL_FILE_NAME = 'cred_file' 112 constants.GCP_ACCESS_TOKEN = 'access_token' 113 tmp_folder = tempfile.mkdtemp() 114 not_upload_file = os.path.join(tmp_folder, 115 constants.DO_NOT_UPLOAD) 116 fake_cred_file = os.path.join(tmp_folder, 117 constants.CREDENTIAL_FILE_NAME) 118 119 # mock cred file not exist, and not_upload file exists. 120 if os.path.exists(fake_cred_file): 121 os.remove(fake_cred_file) 122 if os.path.exists(not_upload_file): 123 os.remove(not_upload_file) 124 Path(not_upload_file).touch() 125 126 atest_gcp_utils.fetch_credential(tmp_folder, dict()) 127 self.assertEqual(0, mock_get_credential_with_auth_flow.call_count) 128 self.assertTrue(os.path.exists(not_upload_file)) 129 130 # mock cred file exists, and not_upload_file not exist. 131 if os.path.exists(fake_cred_file): 132 os.remove(fake_cred_file) 133 Path(fake_cred_file).touch() 134 if os.path.exists(not_upload_file): 135 os.remove(not_upload_file) 136 137 atest_gcp_utils.fetch_credential(tmp_folder, dict()) 138 self.assertEqual(1, mock_get_credential_with_auth_flow.call_count) 139 self.assertFalse(os.path.exists(not_upload_file)) 140