From 00a527999f7ab0f06a48199f9302753800405093 Mon Sep 17 00:00:00 2001
From: James Saryerwinnie <js@jamesls.com>
Date: Fri, 18 May 2018 22:01:20 -0700
Subject: [PATCH] Add option for providing function name to logs cmd

Closes #841.
---
 chalice/cli/__init__.py          | 11 +++++----
 tests/functional/cli/test_cli.py | 38 +++++++++++++++++++++++++++++++-
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/chalice/cli/__init__.py b/chalice/cli/__init__.py
index 3eebad6ab..9c5c79e4c 100644
--- a/chalice/cli/__init__.py
+++ b/chalice/cli/__init__.py
@@ -189,16 +189,19 @@ def delete(ctx, profile, stage):
               default=False,
               help='Controls whether or not lambda log messages are included.')
 @click.option('--stage', default=DEFAULT_STAGE_NAME)
+@click.option('-n', '--name',
+              help='The name of the lambda function to retrieve logs from.',
+              default='api_handler')
 @click.option('--profile', help='The profile to use for fetching logs.')
 @click.pass_context
-def logs(ctx, num_entries, include_lambda_messages, stage, profile):
-    # type: (click.Context, int, bool, str, str) -> None
+def logs(ctx, num_entries, include_lambda_messages, stage, name, profile):
+    # type: (click.Context, int, bool, str, str, str) -> None
     factory = ctx.obj['factory']  # type: CLIFactory
     factory.profile = profile
     config = factory.create_config_obj(stage, False)
     deployed = config.deployed_resources(stage)
-    if deployed is not None and 'api_handler' in deployed.resource_names():
-        lambda_arn = deployed.resource_values('api_handler')['lambda_arn']
+    if name in deployed.resource_names():
+        lambda_arn = deployed.resource_values(name)['lambda_arn']
         session = factory.create_botocore_session()
         retriever = factory.create_log_retriever(
             session, lambda_arn)
diff --git a/tests/functional/cli/test_cli.py b/tests/functional/cli/test_cli.py
index a94b555c5..29cf9cc54 100644
--- a/tests/functional/cli/test_cli.py
+++ b/tests/functional/cli/test_cli.py
@@ -9,9 +9,18 @@
 
 from chalice import cli
 from chalice.cli import factory
-from chalice.config import Config
+from chalice.config import Config, DeployedResources
 from chalice.utils import record_deployed_values
 from chalice.constants import DEFAULT_APIGATEWAY_STAGE_NAME
+from chalice.logs import LogRetriever
+
+
+class FakeConfig(object):
+    def __init__(self, deployed_resources):
+        self._deployed_resources = deployed_resources
+
+    def deployed_resources(self, chalice_stage_name):
+        return self._deployed_resources
 
 
 @pytest.fixture
@@ -318,3 +327,30 @@ def test_can_specify_profile_for_logs(runner, mock_cli_factory):
         )
         assert result.exit_code == 0
         assert mock_cli_factory.profile == 'my-profile'
+
+
+def test_can_provide_lambda_name_for_logs(runner, mock_cli_factory):
+    deployed_resources = DeployedResources({
+        "resources": [
+            {"name": "foo",
+             "lambda_arn": "arn:aws:lambda::app-dev-foo",
+             "resource_type": "lambda_function"}]
+    })
+    mock_cli_factory.create_config_obj.return_value = FakeConfig(
+        deployed_resources)
+    log_retriever = mock.Mock(spec=LogRetriever)
+    log_retriever.retrieve_logs.return_value = []
+    mock_cli_factory.create_log_retriever.return_value = log_retriever
+    with runner.isolated_filesystem():
+        cli.create_new_project_skeleton('testproject')
+        os.chdir('testproject')
+        result = _run_cli_command(
+            runner, cli.logs, ['--name', 'foo'],
+            cli_factory=mock_cli_factory
+        )
+        assert result.exit_code == 0
+    log_retriever.retrieve_logs.assert_called_with(
+        include_lambda_messages=False, max_entries=None)
+    mock_cli_factory.create_log_retriever.assert_called_with(
+        mock.sentinel.Session, 'arn:aws:lambda::app-dev-foo'
+    )