|
| 1 | +import csv |
| 2 | +import os |
| 3 | +import tempfile |
| 4 | +import unittest |
| 5 | +import warnings |
| 6 | + |
| 7 | +from azure.storage.blob import BlobServiceClient |
| 8 | +from requests import exceptions |
| 9 | +from kbcstorage.buckets import Buckets |
| 10 | +from kbcstorage.jobs import Jobs |
| 11 | +from kbcstorage.files import Files |
| 12 | +from kbcstorage.tables import Tables |
| 13 | +from kbcstorage.workspaces import Workspaces |
| 14 | + |
| 15 | + |
| 16 | +class TestWorkspaces(unittest.TestCase): |
| 17 | + def setUp(self): |
| 18 | + self.workspaces = Workspaces(os.getenv('KBC_TEST_API_URL'), os.getenv('KBC_TEST_TOKEN')) |
| 19 | + self.buckets = Buckets(os.getenv('KBC_TEST_API_URL'), os.getenv('KBC_TEST_TOKEN')) |
| 20 | + self.jobs = Jobs(os.getenv('KBC_TEST_API_URL'), os.getenv('KBC_TEST_TOKEN')) |
| 21 | + self.tables = Tables(os.getenv('KBC_TEST_API_URL'), os.getenv('KBC_TEST_TOKEN')) |
| 22 | + self.files = Files(os.getenv('KBC_TEST_API_URL'), os.getenv('KBC_TEST_TOKEN')) |
| 23 | + try: |
| 24 | + file_list = self.files.list(tags=['sapi-client-python-tests']) |
| 25 | + for file in file_list: |
| 26 | + self.files.delete(file['id']) |
| 27 | + except exceptions.HTTPError as e: |
| 28 | + if e.response.status_code != 404: |
| 29 | + raise |
| 30 | + try: |
| 31 | + self.buckets.delete('in.c-py-test-buckets', force=True) |
| 32 | + except exceptions.HTTPError as e: |
| 33 | + if e.response.status_code != 404: |
| 34 | + raise |
| 35 | + # https://github.com/boto/boto3/issues/454 |
| 36 | + warnings.simplefilter("ignore", ResourceWarning) |
| 37 | + |
| 38 | + def tearDown(self): |
| 39 | + try: |
| 40 | + if hasattr(self, 'workspace_id'): |
| 41 | + self.workspaces.delete(self.workspace_id) |
| 42 | + except exceptions.HTTPError as e: |
| 43 | + if e.response.status_code != 404: |
| 44 | + raise |
| 45 | + try: |
| 46 | + self.buckets.delete('in.c-py-test-tables', force=True) |
| 47 | + except exceptions.HTTPError as e: |
| 48 | + if e.response.status_code != 404: |
| 49 | + raise |
| 50 | + |
| 51 | + def test_create_workspace(self): |
| 52 | + workspace = self.workspaces.create() |
| 53 | + self.workspace_id = workspace['id'] |
| 54 | + with self.subTest(): |
| 55 | + self.assertTrue('id' in workspace) |
| 56 | + with self.subTest(): |
| 57 | + self.assertTrue('type' in workspace) |
| 58 | + self.assertTrue(workspace['type'] in ['table', 'file']) |
| 59 | + with self.subTest(): |
| 60 | + self.assertTrue('name' in workspace) |
| 61 | + with self.subTest(): |
| 62 | + self.assertTrue('component' in workspace) |
| 63 | + with self.subTest(): |
| 64 | + self.assertTrue('configurationId' in workspace) |
| 65 | + with self.subTest(): |
| 66 | + self.assertTrue('created' in workspace) |
| 67 | + with self.subTest(): |
| 68 | + self.assertTrue('connection' in workspace) |
| 69 | + with self.subTest(): |
| 70 | + self.assertTrue('backend' in workspace['connection']) |
| 71 | + with self.subTest(): |
| 72 | + self.assertTrue('creatorToken' in workspace) |
| 73 | + |
| 74 | + def test_load_tables_to_workspace(self): |
| 75 | + bucket_id = self.buckets.create('py-test-tables')['id'] |
| 76 | + table1_id = self.__create_table(bucket_id, 'test-table-1', {'col1': 'ping', 'col2': 'pong'}) |
| 77 | + table2_id = self.__create_table(bucket_id, 'test-table-2', {'col1': 'king', 'col2': 'kong'}) |
| 78 | + workspace = self.workspaces.create() |
| 79 | + self.workspace_id = workspace['id'] |
| 80 | + job = self.workspaces.load_tables( |
| 81 | + workspace['id'], |
| 82 | + {table1_id: 'destination_1', table2_id: 'destination_2'} |
| 83 | + ) |
| 84 | + self.jobs.block_until_completed(job['id']) |
| 85 | + |
| 86 | + job = self.tables.create_raw( |
| 87 | + bucket_id, |
| 88 | + 'back-and-forth-table', |
| 89 | + data_workspace_id=workspace['id'], |
| 90 | + data_table_name='destination_1' |
| 91 | + ) |
| 92 | + self.jobs.block_until_completed(job['id']) |
| 93 | + |
| 94 | + new_table = self.tables.detail(bucket_id + '.back-and-forth-table') |
| 95 | + self.assertEqual('back-and-forth-table', new_table['name']) |
| 96 | + |
| 97 | + # test load files into an abs workspace |
| 98 | + def test_load_files_to_workspace(self): |
| 99 | + if (os.getenv('SKIP_ABS_TEST')): |
| 100 | + self.skipTest('Skipping ABS test because env var SKIP_ABS_TESTS was set') |
| 101 | + # put a test file to storage |
| 102 | + file, path = tempfile.mkstemp(prefix='sapi-test') |
| 103 | + os.write(file, bytes('fooBar', 'utf-8')) |
| 104 | + os.close(file) |
| 105 | + file_id = self.files.upload_file(path, tags=['sapi-client-python-tests', 'file1']) |
| 106 | + |
| 107 | + # create a workspace and load the file to it |
| 108 | + workspace = self.workspaces.create('abs') |
| 109 | + self.workspace_id = workspace['id'] |
| 110 | + job = self.workspaces.load_files( |
| 111 | + workspace['id'], |
| 112 | + {'tags': ['sapi-client-python-tests'], 'destination': 'data/in/files'} |
| 113 | + ) |
| 114 | + self.jobs.block_until_completed(job['id']) |
| 115 | + |
| 116 | + # assert that the file was loaded to the workspace |
| 117 | + blob_service_client = BlobServiceClient.from_connection_string(workspace['connection']['connectionString']) |
| 118 | + blob_client = blob_service_client.get_blob_client( |
| 119 | + container=workspace['connection']['container'], |
| 120 | + blob='data/in/files/%s' % str(file_id) |
| 121 | + ) |
| 122 | + self.assertEqual('fooBar', blob_client.download_blob().readall().decode('utf-8')) |
| 123 | + |
| 124 | + def __create_table(self, bucket_id, table_name, row): |
| 125 | + file, path = tempfile.mkstemp(prefix='sapi-test') |
| 126 | + with open(path, 'w') as csv_file: |
| 127 | + writer = csv.DictWriter(csv_file, fieldnames=['col1', 'col2'], |
| 128 | + lineterminator='\n', delimiter=',', quotechar='"') |
| 129 | + writer.writeheader() |
| 130 | + writer.writerow(row) |
| 131 | + return self.tables.create(name=table_name, file_path=path, bucket_id=bucket_id) |
0 commit comments