|
2 | 2 | import sys
|
3 | 3 |
|
4 | 4 | from gcloud import datastore
|
| 5 | +from gcloud import storage |
5 | 6 |
|
6 | 7 |
|
7 | 8 | # Defaults from shell environ. May be None.
|
| 9 | +PROJECT_ID = os.getenv('GCLOUD_TESTS_PROJECT_ID') |
8 | 10 | DATASET_ID = os.getenv('GCLOUD_TESTS_DATASET_ID')
|
9 | 11 | CLIENT_EMAIL = os.getenv('GCLOUD_TESTS_CLIENT_EMAIL')
|
10 | 12 | KEY_FILENAME = os.getenv('GCLOUD_TESTS_KEY_FILE')
|
11 |
| -DATASETS = {} |
| 13 | +CACHED_RETURN_VALS = {} |
12 | 14 |
|
13 | 15 | ENVIRON_ERROR_MSG = """\
|
14 | 16 | To run the regression tests, you need to set some environment variables.
|
15 | 17 | Please check the Contributing guide for instructions.
|
16 | 18 | """
|
17 | 19 |
|
18 | 20 |
|
19 |
| -def get_environ(): |
20 |
| - if DATASET_ID is None or CLIENT_EMAIL is None or KEY_FILENAME is None: |
21 |
| - print >> sys.stderr, ENVIRON_ERROR_MSG |
22 |
| - sys.exit(1) |
| 21 | +def get_environ(require_datastore=False, require_storage=False): |
| 22 | + if require_datastore: |
| 23 | + if DATASET_ID is None or CLIENT_EMAIL is None or KEY_FILENAME is None: |
| 24 | + print >> sys.stderr, ENVIRON_ERROR_MSG |
| 25 | + sys.exit(1) |
| 26 | + |
| 27 | + if require_storage: |
| 28 | + if PROJECT_ID is None or CLIENT_EMAIL is None or KEY_FILENAME is None: |
| 29 | + print >> sys.stderr, ENVIRON_ERROR_MSG |
| 30 | + sys.exit(1) |
23 | 31 |
|
24 | 32 | return {
|
| 33 | + 'project_id': PROJECT_ID, |
25 | 34 | 'dataset_id': DATASET_ID,
|
26 | 35 | 'client_email': CLIENT_EMAIL,
|
27 | 36 | 'key_filename': KEY_FILENAME,
|
28 | 37 | }
|
29 | 38 |
|
30 | 39 |
|
31 | 40 | def get_dataset():
|
32 |
| - environ = get_environ() |
| 41 | + environ = get_environ(require_datastore=True) |
33 | 42 | get_dataset_args = (environ['dataset_id'], environ['client_email'],
|
34 | 43 | environ['key_filename'])
|
35 |
| - if get_dataset_args not in DATASETS: |
| 44 | + key = ('get_dataset', get_dataset_args) |
| 45 | + if key not in CACHED_RETURN_VALS: |
| 46 | + # Cache return value for the environment. |
| 47 | + CACHED_RETURN_VALS[key] = datastore.get_dataset(*get_dataset_args) |
| 48 | + return CACHED_RETURN_VALS[key] |
| 49 | + |
| 50 | + |
| 51 | +def get_storage_connection(): |
| 52 | + environ = get_environ(require_storage=True) |
| 53 | + get_connection_args = (environ['project_id'], environ['client_email'], |
| 54 | + environ['key_filename']) |
| 55 | + key = ('get_storage_connection', get_connection_args) |
| 56 | + if key not in CACHED_RETURN_VALS: |
36 | 57 | # Cache return value for the environment.
|
37 |
| - DATASETS[get_dataset_args] = datastore.get_dataset(*get_dataset_args) |
38 |
| - return DATASETS[get_dataset_args] |
| 58 | + CACHED_RETURN_VALS[key] = storage.get_connection(*get_connection_args) |
| 59 | + return CACHED_RETURN_VALS[key] |
0 commit comments