diff --git a/src/functions_framework/_cli.py b/src/functions_framework/_cli.py index 48455ea6..aed69116 100644 --- a/src/functions_framework/_cli.py +++ b/src/functions_framework/_cli.py @@ -32,13 +32,26 @@ @click.option("--host", envvar="HOST", type=click.STRING, default="0.0.0.0") @click.option("--port", envvar="PORT", type=click.INT, default=8080) @click.option("--debug", envvar="DEBUG", is_flag=True) +@click.option( + "--env", + "env_vars", + multiple=True, + metavar="KEY=VALUE", + help="Set a runtime environment variable for local execution.", +) @click.option( "--asgi", envvar="FUNCTION_USE_ASGI", is_flag=True, help="Use ASGI server for function execution", ) -def _cli(target, source, signature_type, host, port, debug, asgi): +def _cli(target, source, signature_type, host, port, debug, env_vars, asgi): + for env_var in env_vars: + key, separator, value = env_var.partition("=") + if not key or not separator: + raise click.BadParameter("must be in KEY=VALUE format", param_hint="--env") + os.environ[key] = value + if asgi: from functions_framework.aio import create_asgi_app diff --git a/tests/test_cli.py b/tests/test_cli.py index 75c93f20..edc09462 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -130,6 +130,34 @@ def test_cli(monkeypatch, args, env, create_app_calls, run_calls): assert wsgi_server.run.calls == run_calls +def test_cli_sets_runtime_env(monkeypatch): + wsgi_server = pretend.stub(run=pretend.call_recorder(lambda *a, **kw: None)) + wsgi_app = pretend.stub(run=pretend.call_recorder(lambda *a, **kw: None)) + create_app = pretend.call_recorder(lambda *a, **kw: wsgi_app) + monkeypatch.setattr(functions_framework._cli, "create_app", create_app) + create_server = pretend.call_recorder(lambda *a, **kw: wsgi_server) + monkeypatch.setattr(functions_framework._cli, "create_server", create_server) + monkeypatch.delenv("LOCAL_ONLY", raising=False) + monkeypatch.delenv("EMPTY_VALUE", raising=False) + + runner = CliRunner() + result = runner.invoke( + _cli, ["--target", "foo", "--env", "LOCAL_ONLY=1", "--env", "EMPTY_VALUE="] + ) + + assert result.exit_code == 0 + assert os.environ["LOCAL_ONLY"] == "1" + assert os.environ["EMPTY_VALUE"] == "" + + +def test_cli_rejects_invalid_runtime_env(): + runner = CliRunner() + result = runner.invoke(_cli, ["--target", "foo", "--env", "LOCAL_ONLY"]) + + assert result.exit_code == 2 + assert "must be in KEY=VALUE format" in result.output + + def test_asgi_cli(monkeypatch): asgi_server = pretend.stub(run=pretend.call_recorder(lambda *a, **kw: None)) asgi_app = pretend.stub()