import os import requests import json import sys import time class ImportHelper: @classmethod def import_test_dataset(cls, username, api_key, host, schema): requests.packages.urllib3.disable_warnings() url = "{0}://{1}.{2}/api/v1/imports/"\ "?type_guessing=false&privacy=public&api_key={3}".format( schema, username, host, api_key) dataset = { 'file': open('fixtures/geocoder_api_test_dataset.csv', 'rb')} response = requests.post(url, files=dataset) response.raise_for_status() response_json = json.loads(response.text) if not response_json['success']: print("Error importing the test dataset: {0}".format(response.text)) sys.exit(1) while(True): table_name = ImportHelper.get_imported_table_name( username, host, api_key, response_json['item_queue_id'], schema ) if table_name: return table_name else: time.sleep(5) @classmethod def get_imported_table_name(cls, username, host, api_key, import_id, schema): requests.packages.urllib3.disable_warnings() import_url = "{0}://{1}.{2}/api/v1/imports/{3}?api_key={4}".format( schema, username, host, import_id, api_key) import_data_response = requests.get(import_url) if import_data_response.status_code != 200: print("Error getting the table name from " \ "the import data: {0}".format(import_data_response.text)) sys.exit(1) import_data_json = json.loads(import_data_response.text) return import_data_json['table_name'] @classmethod def clean_test_dataset(cls, username, api_key, table_name, host, schema): requests.packages.urllib3.disable_warnings() url = "{0}://{1}.{2}/api/v2/sql?q=drop table {3}&api_key={4}".format( schema, username, host, table_name, api_key ) response = requests.get(url) if response.status_code != 200: print("Error cleaning the test dataset: {0}".format(response.text)) sys.exit(1)