From d249a35c41b702a405ffb42c7f13739af8503956 Mon Sep 17 00:00:00 2001 From: "Philipp v. K" Date: Mon, 22 Aug 2022 08:24:09 +0200 Subject: [PATCH] [Tests] [TVMC] add tests for filter_tasks utility --- tests/python/driver/tvmc/test_autotuner.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/python/driver/tvmc/test_autotuner.py b/tests/python/driver/tvmc/test_autotuner.py index 66017823a6699..5c69fa65bdd24 100644 --- a/tests/python/driver/tvmc/test_autotuner.py +++ b/tests/python/driver/tvmc/test_autotuner.py @@ -24,6 +24,7 @@ from tvm import autotvm from tvm.driver import tvmc +from tvm.driver.tvmc.autotuner import filter_tasks def _get_tasks(model): @@ -182,3 +183,23 @@ def test_tune_rpc_tracker_parsing(mock_load_model, mock_tune_model, mock_auto_sc assert "10.0.0.1" == kwargs["hostname"] assert "port" in kwargs assert 9999 == kwargs["port"] + + +def test_filter_tasks_valid(): + filter_tasks(list(range(10)), "list") == ([], True) + filter_tasks(list(range(10)), "help") == ([], True) + filter_tasks(list(range(10)), "all") == ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], False) + filter_tasks(list(range(10)), "5") == ([5], False) + filter_tasks(list(range(10)), "1-5") == ([1, 2, 3, 4, 5], False) + filter_tasks(list(range(10)), "-5") == ([0, 1, 2, 3, 4, 5], False) + filter_tasks(list(range(10)), "6-") == ([6, 7, 8, 9], False) + filter_tasks(list(range(10)), "0,1-3,all") == ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], False) + filter_tasks(list(range(10)), "0,4-5,9,list") == ([0, 4, 5, 9], True) + + +@pytest.mark.xfail +def test_filter_tasks_invalid(): + filter_tasks(list(range(10)), "10") + filter_tasks(list(range(10)), "5,10") + filter_tasks(list(range(10)), "1-10") + filter_tasks(list(range(10)), "-10")