Repository: PrefectHQ/prefect

Issue:
Add CLI flags for all Fargate Agent keyword arguments
As another interface for setting the various config settings.


Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -20,7 +20,7 @@
     "dropbox": ["dropbox ~= 9.0"],
     "google": [
         "google-cloud-bigquery >= 1.6.0, < 2.0",
-        "google-cloud-storage >= 1.13, < 2.0",
+        "google-cloud-storage >= 1.13, < 1.23.0",
     ],
     "kubernetes": ["kubernetes >= 9.0.0a1, < 10.0", "dask-kubernetes >= 0.8.0"],
     "rss": ["feedparser >= 5.0.1, < 6.0"],
diff --git a/src/prefect/cli/agent.py b/src/prefect/cli/agent.py
--- a/src/prefect/cli/agent.py
+++ b/src/prefect/cli/agent.py
@@ -1,6 +1,6 @@
 import click
 
-from prefect import config, context
+from prefect import config
 from prefect.utilities.configuration import set_temporary_config
 from prefect.utilities.serialization import from_qualified_name
 
@@ -42,7 +42,10 @@ def agent():
     pass
 
 
-@agent.command(hidden=True)
+@agent.command(
+    hidden=True,
+    context_settings=dict(ignore_unknown_options=True, allow_extra_args=True,),
+)
 @click.argument("agent-option", default="local")
 @click.option(
     "--token", "-t", required=False, help="A Prefect Cloud API token.", hidden=True
@@ -67,7 +70,8 @@ def agent():
 )
 @click.option("--no-pull", is_flag=True, help="Pull images flag.", hidden=True)
 @click.option("--base-url", "-b", help="Docker daemon base URL.", hidden=True)
-def start(agent_option, token, name, verbose, label, no_pull, base_url):
+@click.pass_context
+def start(ctx, agent_option, token, name, verbose, label, no_pull, base_url):
     """
     Start an agent.
 
@@ -90,7 +94,19 @@ def start(agent_option, token, name, verbose, label, no_pull, base_url):
         --base-url, -b  TEXT    A Docker daemon host URL for a LocalAgent
         --no-pull               Pull images for a LocalAgent
                                 Defaults to pulling if not provided
+
+    \b
+    Fargate Agent Options:
+        Any of the configuration options outlined in the docs can be provided here
+        https://docs.prefect.io/cloud/agent/fargate.html#configuration
     """
+
+    # Split context
+    kwargs = dict()
+    for item in ctx.args:
+        item = item.replace("--", "")
+        kwargs.update([item.split("=")])
+
     tmp_config = {"cloud.agent.auth_token": token or config.cloud.agent.auth_token}
     if verbose:
         tmp_config["cloud.agent.level"] = "DEBUG"
@@ -102,7 +118,17 @@ def start(agent_option, token, name, verbose, label, no_pull, base_url):
             click.secho("{} is not a valid agent".format(agent_option), fg="red")
             return
 
-        with context(no_pull=no_pull, base_url=base_url):
+        _agent = from_qualified_name(retrieved_agent)
+
+        if agent_option == "local":
+            from_qualified_name(retrieved_agent)(
+                name=name, labels=list(label), base_url=base_url, no_pull=no_pull,
+            ).start()
+        elif agent_option == "fargate":
+            from_qualified_name(retrieved_agent)(
+                name=name, labels=list(label), **kwargs
+            ).start()
+        else:
             from_qualified_name(retrieved_agent)(name=name, labels=list(label)).start()
 
 


Test patch:
diff --git a/tests/cli/test_agent.py b/tests/cli/test_agent.py
--- a/tests/cli/test_agent.py
+++ b/tests/cli/test_agent.py
@@ -1,9 +1,14 @@
 from unittest.mock import MagicMock
 
 from click.testing import CliRunner
+import pytest
 
 from prefect.cli.agent import agent
 
+pytest.importorskip("boto3")
+pytest.importorskip("botocore")
+pytest.importorskip("kubernetes")
+
 
 def test_agent_init():
     runner = CliRunner()
@@ -55,6 +60,88 @@ def test_agent_start_verbose(monkeypatch, runner_token):
     assert result.exit_code == 0
 
 
+def test_agent_start_local(monkeypatch, runner_token):
+    start = MagicMock()
+    monkeypatch.setattr("prefect.agent.local.LocalAgent.start", start)
+
+    docker_client = MagicMock()
+    monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", docker_client)
+
+    runner = CliRunner()
+    result = runner.invoke(agent, ["start", "local"])
+    assert result.exit_code == 0
+
+
+def test_agent_start_kubernetes(monkeypatch, runner_token):
+    start = MagicMock()
+    monkeypatch.setattr("prefect.agent.kubernetes.KubernetesAgent.start", start)
+
+    k8s_config = MagicMock()
+    monkeypatch.setattr("kubernetes.config", k8s_config)
+
+    runner = CliRunner()
+    result = runner.invoke(agent, ["start", "kubernetes"])
+    assert result.exit_code == 0
+
+
+def test_agent_start_kubernetes_kwargs_ignored(monkeypatch, runner_token):
+    start = MagicMock()
+    monkeypatch.setattr("prefect.agent.kubernetes.KubernetesAgent.start", start)
+
+    k8s_config = MagicMock()
+    monkeypatch.setattr("kubernetes.config", k8s_config)
+
+    runner = CliRunner()
+    result = runner.invoke(agent, ["start", "kubernetes", "test_kwarg=ignored"])
+    assert result.exit_code == 0
+
+
+def test_agent_start_fargate(monkeypatch, runner_token):
+    start = MagicMock()
+    monkeypatch.setattr("prefect.agent.fargate.FargateAgent.start", start)
+
+    boto3_client = MagicMock()
+    monkeypatch.setattr("boto3.client", boto3_client)
+
+    runner = CliRunner()
+    result = runner.invoke(agent, ["start", "fargate"])
+    assert result.exit_code == 0
+
+
+def test_agent_start_fargate_kwargs(monkeypatch, runner_token):
+    start = MagicMock()
+    monkeypatch.setattr("prefect.agent.fargate.FargateAgent.start", start)
+
+    boto3_client = MagicMock()
+    monkeypatch.setattr("boto3.client", boto3_client)
+
+    runner = CliRunner()
+    result = runner.invoke(agent, ["start", "fargate", "taskRoleArn=test"])
+    assert result.exit_code == 0
+
+
+def test_agent_start_fargate_kwargs_received(monkeypatch, runner_token):
+    start = MagicMock()
+    monkeypatch.setattr("prefect.agent.fargate.FargateAgent.start", start)
+
+    fargate_agent = MagicMock()
+    monkeypatch.setattr("prefect.agent.fargate.FargateAgent", fargate_agent)
+
+    boto3_client = MagicMock()
+    monkeypatch.setattr("boto3.client", boto3_client)
+
+    runner = CliRunner()
+    result = runner.invoke(
+        agent, ["start", "fargate", "taskRoleArn=arn", "--volumes=vol"]
+    )
+    assert result.exit_code == 0
+
+    assert fargate_agent.called
+    fargate_agent.assert_called_with(
+        labels=[], name=None, taskRoleArn="arn", volumes="vol"
+    )
+
+
 def test_agent_start_name(monkeypatch, runner_token):
     start = MagicMock()
     monkeypatch.setattr("prefect.agent.local.LocalAgent.start", start)
@@ -67,7 +154,7 @@ def test_agent_start_name(monkeypatch, runner_token):
     assert result.exit_code == 0
 
 
-def test_agent_start_local_context_vars(monkeypatch, runner_token):
+def test_agent_start_local_vars(monkeypatch, runner_token):
     start = MagicMock()
     monkeypatch.setattr("prefect.agent.local.LocalAgent.start", start)