Merge pull request #2668 from cortesi/commandopts

Command improvements

2 files changed, 22 insertions, 6 deletions

diff --git a/test/mitmproxy/addons/test_core.py b/test/mitmproxy/addons/test_core.py
index c132d80a..5aa4ef37 100644
--- a/test/mitmproxy/addons/test_core.py
+++ b/test/mitmproxy/addons/test_core.py
@@ -69,9 +69,6 @@ def test_flow_set():
         f = tflow.tflow(resp=True)
         assert sa.flow_set_options()
 
-        with pytest.raises(exceptions.CommandError):
-            sa.flow_set([f], "flibble", "post")
-
         assert f.request.method != "post"
         sa.flow_set([f], "method", "post")
         assert f.request.method == "POST"
@@ -126,9 +123,6 @@ def test_encoding():
         sa.encode_toggle([f], "request")
         assert "content-encoding" not in f.request.headers
 
-        with pytest.raises(exceptions.CommandError):
-            sa.encode([f], "request", "invalid")
-
 
 def test_options(tmpdir):
     p = str(tmpdir.join("path"))
diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py
index 43b97742..cb9dc4ed 100644
--- a/test/mitmproxy/test_command.py
+++ b/test/mitmproxy/test_command.py
@@ -8,6 +8,10 @@ import io
 import pytest
 
 
+TChoice = typing.NewType("TChoice", command.Choice)
+TChoice.options_command = "choices"
+
+
 class TAddon:
     def cmd1(self, foo: str) -> str:
         """cmd1 help"""
@@ -25,6 +29,12 @@ class TAddon:
     def varargs(self, one: str, *var: str) -> typing.Sequence[str]:
         return list(var)
 
+    def choices(self) -> typing.Sequence[str]:
+        return ["one", "two", "three"]
+
+    def choose(self, arg: TChoice) -> typing.Sequence[str]:  # type: ignore
+        return ["one", "two", "three"]
+
 
 class TestCommand:
     def test_varargs(self):
@@ -86,6 +96,8 @@ def test_typename():
     assert command.typename(flow.Flow, False) == "flow"
     assert command.typename(typing.Sequence[str], False) == "[str]"
 
+    assert command.typename(TChoice, False) == "choice"
+
 
 class DummyConsole:
     @command.command("view.resolve")
@@ -134,6 +146,16 @@ def test_parsearg():
             tctx.master.commands, "foo, bar", typing.Sequence[str]
         ) == ["foo", "bar"]
 
+        a = TAddon()
+        tctx.master.commands.add("choices", a.choices)
+        assert command.parsearg(
+            tctx.master.commands, "one", TChoice,
+        ) == "one"
+        with pytest.raises(exceptions.CommandError):
+            assert command.parsearg(
+                tctx.master.commands, "invalid", TChoice,
+            )
+
 
 class TDec:
     @command.command("cmd1")