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 __future__ import absolute_import
from __future__ import print_function
from datetime import datetime
import os
import re

from dateutil import tz
from googleapiclient import discovery
import oauth2client
import oauth2client.file
import oauth2client.tools
from six.moves import range

# 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'

REQUEST_RETRIES = 3


class GoogleDriveDownloadException(Exception):
  """Franky class for exceptions related to Google Sheet downloading errors."""
  pass


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, folder):
  """Download Franky test cases spreadsheet as CSV file into folder.

  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
    folder: (str) A path where GoogleSheet will be downloaded.

  Returns:
    A list of full file paths of all the files that are downloaded.

  Raises:
    GoogleDriveDownloadError: if workbook has not been downloaded.
  """
  credentials = get_credentials(flags)
  drive_service = discovery.build(DRIVE_API_NAME, DRIVE_API_VERSION,
                                  credentials=credentials,
                                  cache_discovery=False)
  test_suite = []
  for workbook_id in workbooks_list:
    for attempt in range(0, REQUEST_RETRIES):
      response, content = drive_service._http.request(
          SPREADSHEET_CSV_URL_TEMPLATE.format(key=spreadsheet_key,
                                              id=workbook_id))
      if response.status == 200:
        test_suite.append(SaveResponseToCSV(response, content, folder))
        break
      else:
        print(('#%d attempt to download. '
               'An error occurred with the status code: %s' % (
                   attempt + 1, response.status)))
    if response.status != 200:
      raise GoogleDriveDownloadException(
          'Did not download any spreadsheet for [%s, %s]!' % (
              spreadsheet_key, ','.join(workbooks_list)))
  return test_suite


def SaveResponseToCSV(response, content, folder):
  """Creates csv file based on GoogleDrive downloading response.

  Args:
    response: A response of downloading request,
    content: A content of response.
    folder: A folder where CSV file with content will be stored.

  Returns:
    Full path of new created CSV file.
  """
  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' % (
      folder, current_time_stamp, csv_file_name)
  with open(csv_file_full_path, 'w') as csv_file:
    csv_file.write(content.decode('utf-8'))
  return csv_file_full_path