commands: add a command.command decorator

Use this for our built-ins and the console commands.

2 files changed, 56 insertions, 16 deletions

diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py
index 034182a6..678bc1b7 100644
--- a/test/mitmproxy/test_addonmanager.py
+++ b/test/mitmproxy/test_addonmanager.py
@@ -4,6 +4,7 @@ from mitmproxy import addons
 from mitmproxy import addonmanager
 from mitmproxy import exceptions
 from mitmproxy import options
+from mitmproxy import command
 from mitmproxy import master
 from mitmproxy import proxy
 from mitmproxy.test import taddons
@@ -18,6 +19,10 @@ class TAddon:
         if addons:
             self.addons = addons
 
+    @command.command("test.command")
+    def testcommand(self) -> str:
+        return "here"
+
     def __repr__(self):
         return "Addon(%s)" % self.name
 
@@ -38,6 +43,12 @@ class AOption:
         l.add_option("custom_option", bool, False, "help")
 
 
+def test_command():
+    with taddons.context() as tctx:
+        tctx.master.addons.add(TAddon("test"))
+        assert tctx.master.commands.call("test.command") == "here"
+
+
 def test_halt():
     o = options.Options()
     m = master.Master(o, proxy.DummyServer(o))
diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py
index 16bc8a90..64928dbf 100644
--- a/test/mitmproxy/test_command.py
+++ b/test/mitmproxy/test_command.py
@@ -39,22 +39,21 @@ class TestCommand:
 
 
 def test_simple():
-    o = options.Options()
-    m = master.Master(o, proxy.DummyServer(o))
-    c = command.CommandManager(m)
-    a = TAddon()
-    c.add("one.two", a.cmd1)
-    assert c.commands["one.two"].help == "cmd1 help"
-    assert(c.call("one.two foo") == "ret foo")
-    with pytest.raises(exceptions.CommandError, match="Unknown"):
-        c.call("nonexistent")
-    with pytest.raises(exceptions.CommandError, match="Invalid"):
-        c.call("")
-    with pytest.raises(exceptions.CommandError, match="Usage"):
-        c.call("one.two too many args")
-
-    c.add("empty", a.empty)
-    c.call("empty")
+    with taddons.context() as tctx:
+        c = command.CommandManager(tctx.master)
+        a = TAddon()
+        c.add("one.two", a.cmd1)
+        assert c.commands["one.two"].help == "cmd1 help"
+        assert(c.call("one.two foo") == "ret foo")
+        with pytest.raises(exceptions.CommandError, match="Unknown"):
+            c.call("nonexistent")
+        with pytest.raises(exceptions.CommandError, match="Invalid"):
+            c.call("")
+        with pytest.raises(exceptions.CommandError, match="Usage"):
+            c.call("one.two too many args")
+
+        c.add("empty", a.empty)
+        c.call("empty")
 
 
 def test_typename():
@@ -87,3 +86,33 @@ def test_parsearg():
             command.parsearg(tctx.master.commands, "0", flow.Flow)
         with pytest.raises(exceptions.CommandError):
             command.parsearg(tctx.master.commands, "foo", Exception)
+
+
+class TDec:
+    @command.command("cmd1")
+    def cmd1(self, foo: str) -> str:
+        """cmd1 help"""
+        return "ret " + foo
+
+    @command.command("cmd2")
+    def cmd2(self, foo: str) -> str:
+        return 99
+
+    @command.command("empty")
+    def empty(self) -> None:
+        pass
+
+
+def test_decorator():
+    with taddons.context() as tctx:
+        c = command.CommandManager(tctx.master)
+        a = TDec()
+        c.collect_commands(a)
+        assert "cmd1" in c.commands
+        assert c.call("cmd1 bar") == "ret bar"
+        assert "empty" in c.commands
+        assert c.call("empty") is None
+
+    with taddons.context() as tctx:
+        tctx.master.addons.add(a)
+        assert tctx.master.commands.call("cmd1 bar") == "ret bar"