diff --git a/saws/completer.py b/saws/completer.py index c2d974b4..3a0376dd 100644 --- a/saws/completer.py +++ b/saws/completer.py @@ -88,7 +88,22 @@ def __init__(self, self.config_obj['main'].as_bool('refresh_instance_ids'), self.config_obj['main'].as_bool('refresh_instance_tags'), self.config_obj['main'].as_bool('refresh_bucket_names')) - self.resources.refresh() + self.resource_map = None + self.refresh_resources() + + def create_resource_map(self): + """Creates a mapping of resource keywords and resources to complete. + + Example: + Key: '--instance-ids'. + Value: List of instance ids. + + Args: + * None. + + Returns: + None. + """ self.resource_map = dict(zip([self.resources.INSTANCE_IDS, self.resources.EC2_TAG_KEY, self.resources.EC2_TAG_VALUE, @@ -102,6 +117,19 @@ def __init__(self, self.resources.bucket_names, self.resources.s3_uri_names])) + def refresh_resources(self, force_refresh=False): + """Convenience function to refresh resources for completion. + + Args: + * force_refresh: A boolean determines whether to force a cache + refresh. This value is set to True when the user presses `F5`. + + Returns: + None. + """ + self.resources.refresh(force_refresh) + self.create_resource_map() + def replace_shortcut(self, text): """Replaces matched shortcut commands with their full command. diff --git a/saws/resources.py b/saws/resources.py index d2802646..c251b41c 100644 --- a/saws/resources.py +++ b/saws/resources.py @@ -92,8 +92,8 @@ def __init__(self, self.instance_ids = [] self.instance_tag_keys = set() self.instance_tag_values = set() - self.bucket_names = [] - self.s3_uri_names = [] + self.bucket_names = [] # TODO: Make this 'private' + self.s3_uri_names = [] # TODO: Make this 'private' self.refresh_instance_ids = refresh_instance_ids self.refresh_instance_tags = refresh_instance_tags self.refresh_bucket_names = refresh_bucket_names @@ -215,7 +215,7 @@ def query_bucket_names(self): """ output = self.query_aws(self.QUERY_BUCKET_NAMES_CMD) if output is not None: - self.bucket_names = [] + self.clear_bucket_names() result_list = output.split('\n') for result in result_list: try: @@ -237,6 +237,20 @@ def add_bucket_name(self, bucket_name): self.bucket_names.append(bucket_name) self.s3_uri_names.append(self.S3_URI + '//' + bucket_name) + def clear_bucket_names(self): + """Clears bucket all bucket names. + + Long description. + + Args: + * None. + + Returns: + None. + """ + self.bucket_names = [] + self.s3_uri_names = [] + def refresh_resources_from_file(self, file_path): """Refreshes the AWS resources from data/RESOURCES.txt. @@ -251,7 +265,7 @@ def refresh_resources_from_file(self, file_path): self.instance_ids = [] self.instance_tag_keys = set() self.instance_tag_values = set() - self.bucket_names = [] + self.clear_bucket_names() instance_tag_keys_list = [] instance_tag_values_list = [] for line in fp: diff --git a/tests/test_completer.py b/tests/test_completer.py index be3450db..3f24b413 100644 --- a/tests/test_completer.py +++ b/tests/test_completer.py @@ -24,6 +24,7 @@ from saws.commands import AwsCommands from saws.saws import Saws from saws.resources import AwsResources +from test_resources import ResourcesTest class CompleterTest(unittest.TestCase): @@ -216,3 +217,26 @@ def test_substitutions(self): expected = 'aws ec2 describe-instances --filters "Name=tag-value,Values=prod"' result = self.completer.replace_shortcut(command) assert result == expected + + @mock.patch('saws.resources.print') + def test_refresh_resources(self, mock_print): + NUM_EC2_STATE = 6 + self.completer.resources.RESOURCE_FILE = \ + ResourcesTest.RESOURCES_SAMPLE + self.completer.resource_map = None + self.completer.refresh_resources(force_refresh=False) + mock_print.assert_called_with('Loaded resources from cache') + keys = [self.completer.resources.INSTANCE_IDS, + self.completer.resources.EC2_TAG_KEY, + self.completer.resources.EC2_TAG_VALUE, + self.completer.resources.EC2_STATE, + self.completer.resources.BUCKET, + self.completer.resources.S3_URI] + expected = [ResourcesTest.NUM_INSTANCE_IDS, + ResourcesTest.NUM_INSTANCE_TAG_KEYS, + ResourcesTest.NUM_INSTANCE_TAG_VALUES, + NUM_EC2_STATE, + ResourcesTest.NUM_BUCKET_NAMES, + ResourcesTest.NUM_BUCKET_NAMES] + for i in range(len(keys)): + assert len(self.completer.resource_map[keys[i]]) == expected[i] diff --git a/tests/test_resources.py b/tests/test_resources.py index f4ba6537..917a49e9 100644 --- a/tests/test_resources.py +++ b/tests/test_resources.py @@ -27,13 +27,14 @@ class ResourcesTest(unittest.TestCase): + NUM_INSTANCE_IDS = 7 + NUM_INSTANCE_TAG_KEYS = 3 + NUM_INSTANCE_TAG_VALUES = 6 + NUM_BUCKET_NAMES = 16 + RESOURCES = 'data/RESOURCES.txt' + RESOURCES_SAMPLE = 'data/RESOURCES_SAMPLE.txt' + def setUp(self): - self.NUM_INSTANCE_IDS = 7 - self.NUM_INSTANCE_TAG_KEYS = 3 - self.NUM_INSTANCE_TAG_VALUES = 6 - self.NUM_BUCKET_NAMES = 16 - self.RESOURCES = 'data/RESOURCES.txt' - self.RESOURCES_SAMPLE = 'data/RESOURCES_SAMPLE.txt' self.create_resources() @mock.patch('saws.resources.print') @@ -54,6 +55,8 @@ def test_refresh(self, mock_print): self.NUM_INSTANCE_TAG_VALUES assert len(self.resources.bucket_names) == \ self.NUM_BUCKET_NAMES + assert len(self.resources.s3_uri_names) == \ + self.NUM_BUCKET_NAMES mock_print.assert_called_with('Loaded resources from cache') @mock.patch('saws.resources.subprocess') @@ -106,3 +109,14 @@ def test_query_aws_with_bucket_names(self, mock_subprocess): self.resources.QUERY_BUCKET_NAMES_CMD, universal_newlines=True, shell=True) + + def test_add_and_clear_bucket_name(self): + BUCKET_NAME = 'test_bucket_name' + self.resources.clear_bucket_names() + self.resources.add_bucket_name(BUCKET_NAME) + assert BUCKET_NAME in self.resources.bucket_names + assert str(self.resources.S3_URI + '//' + BUCKET_NAME) in \ + self.resources.s3_uri_names + self.resources.clear_bucket_names() + assert len(self.resources.bucket_names) == 0 + assert len(self.resources.s3_uri_names) == 0