franky/download_testcases_csv.py

# Copyright 2017 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Download Franky test cases from Google Sheets as a CSV file.
"""

from datetime import datetime
from dateutil import tz
import os
import re
import tempfile

from googleapiclient import discovery
import oauth2client
import oauth2client.file
import oauth2client.tools

# Permission required to authorize users to access Google Drive.
DRIVE_SCOPE = 'https://www.googleapis.com/auth/drive.readonly'

# Authoriztion details for accessing Google APIs using OAuth2.0
CLIENT_SECRET_FILE = 'client_secret.json'

# Name of the Google Drive API.
DRIVE_API_NAME = 'drive'

# Version of the Google Drive API.
DRIVE_API_VERSION = 'v3'

# URL for downloading spreadsheet as CSV.
SPREADSHEET_CSV_URL_TEMPLATE = ('https://docs.google.com/spreadsheets/export?'
                                'id={key}&format=csv&gid={id}')

FILE_NAME_PATTERN = r'filename="(.*).csv"'
TIME_STAMP_PATTERN = r'date\': \'(.*) GMT'


def get_credentials(flags):
  """Gets valid user credentials from storage.

  If nothing has been stored, or if the stored credentials are invalid,
  the OAuth2 flow is completed to obtain the new credentials.

  Args:
    flags: The command-line flags required for one-time authorization.

  Returns:
    Credentials, the obtained credential.
  """
  home_dir = os.path.expanduser('~')
  credential_dir = os.path.join(home_dir, '.credentials')
  if not os.path.exists(credential_dir):
    os.makedirs(credential_dir)
  credential_path = os.path.join(credential_dir, 'credentials.json')
  store = oauth2client.file.Storage(credential_path)
  credentials = store.get()
  if not credentials or credentials.invalid:
    flow = oauth2client.client.flow_from_clientsecrets(CLIENT_SECRET_FILE,
                                                       DRIVE_SCOPE)
    credentials = oauth2client.tools.run_flow(flow, store, flags)
  return credentials


def DownloadCSVFromGoogleSheet(flags, spreadsheet_key, workbooks_list):
  """Download Franky test cases spreadsheet as CSV file into temp directory.

  Args:
    flags: The command-line flags required for one-time authorization.
    spreadsheet_key: (str) Id of GoogleSheet document.
    workbooks_list: A list of Workbook IDs. ID is the 10 digits after #gid=" in
                    spreadsheet tab URL

  Returns:
    A list of full file paths of all the files that are downloaded.
  """
  credentials = get_credentials(flags)
  drive_service = discovery.build(DRIVE_API_NAME, DRIVE_API_VERSION,
                                  credentials=credentials,
                                  cache_discovery=False)
  tmp_dir = tempfile.mkdtemp()
  test_suite = []
  for workbook_id in workbooks_list:
    response, content = drive_service._http.request(
        SPREADSHEET_CSV_URL_TEMPLATE.format(key=spreadsheet_key,
                                            id=workbook_id))
    if response.status == 200:
      csv_file_name = re.search(FILE_NAME_PATTERN, str(response)).group(1)
      timestamp = re.search(TIME_STAMP_PATTERN, str(response)).group(1)
      utc = datetime.strptime(timestamp, '%a, %d %b %Y %H:%M:%S')
      utc = utc.replace(tzinfo=tz.tzutc())
      # Convert to local time zone
      current_time_stamp = utc.astimezone(tz.tzlocal())
      csv_file_full_path = '%s/%s %s.csv' % (
          tmp_dir, current_time_stamp, csv_file_name)
      with open(csv_file_full_path, 'w') as csv_file:
        csv_file.write(content)
      test_suite.append(csv_file_full_path)
    else:
      print 'An error occurred with the status code: %s' % response.status
  return test_suite