3 files changed, 126 insertions, 2 deletions
diff --git a/tools/tryselect/test/test_fuzzy.py b/tools/tryselect/test/test_fuzzy.py
index 9ff1b386af..e68365c923 100644
--- a/tools/tryselect/test/test_fuzzy.py
+++ b/tools/tryselect/test/test_fuzzy.py
@@ -98,7 +98,7 @@ def test_query_paths_variants(run_mach, capfd, variant):
     index = output.find(delim)
     result = json.loads(output[index + len(delim) :])
     tasks = result["parameters"]["try_task_config"]["tasks"]
-    assert tasks == expected
+    assert sorted(tasks) == sorted(expected)
 
 
 @pytest.mark.skipif(os.name == "nt", reason="fzf not installed on host")
diff --git a/tools/tryselect/test/test_perf.py b/tools/tryselect/test/test_perf.py
index 0db45df83e..c4e3b6b5fa 100644
--- a/tools/tryselect/test/test_perf.py
+++ b/tools/tryselect/test/test_perf.py
@@ -815,6 +815,34 @@ def test_category_expansion(
 
 
 @pytest.mark.parametrize(
+    "category_options, call_counts",
+    [
+        (
+            {},
+            0,
+        ),
+        (
+            {"non_pgo": True},
+            58,
+        ),
+    ],
+)
+def test_category_expansion_with_non_pgo_flag(category_options, call_counts):
+    PerfParser.categories = TEST_CATEGORIES
+    PerfParser.variants = TEST_VARIANTS
+
+    expanded_cats = PerfParser.get_categories(**category_options)
+
+    non_shippable_count = 0
+    for cat_name in expanded_cats:
+        queries = str(expanded_cats[cat_name].get("queries", None))
+        if "!shippable !nightlyasrelease" in queries and "'shippable" not in queries:
+            non_shippable_count += 1
+
+    assert non_shippable_count == call_counts
+
+
+@pytest.mark.parametrize(
     "options, call_counts, log_ind, expected_log_message",
     [
         (
diff --git a/tools/tryselect/test/test_tasks.py b/tools/tryselect/test/test_tasks.py
index 2e99c72d8b..ecf5849741 100644
--- a/tools/tryselect/test/test_tasks.py
+++ b/tools/tryselect/test/test_tasks.py
@@ -6,7 +6,103 @@ import os
 
 import mozunit
 import pytest
-from tryselect.tasks import cache_key, filter_tasks_by_paths, resolve_tests_by_suite
+from tryselect.tasks import (
+    cache_key,
+    filter_tasks_by_paths,
+    filter_tasks_by_worker_type,
+    resolve_tests_by_suite,
+)
+
+
+class task:
+    def __init__(self, workerType):
+        self.workerType = workerType
+
+    @property
+    def task(self):
+        return {"workerType": self.workerType}
+
+
+@pytest.mark.parametrize(
+    "tasks, params, expected",
+    (
+        pytest.param(
+            {
+                "foobar/xpcshell-1": task("t-unittest-314"),
+                "foobar/mochitest": task("t-unittest-157"),
+                "foobar/xpcshell-gpu": task("t-unittest-314-gpu"),
+                "foobar/xpcshell": task("t-unittest-314"),
+            },
+            {"try_task_config": {"worker-types": ["t-unittest-314"]}},
+            [
+                "foobar/xpcshell-1",
+                "foobar/xpcshell",
+            ],
+            id="single worker",
+        ),
+        pytest.param(
+            {
+                "foobar/xpcshell-1": task("t-unittest-314"),
+                "foobar/mochitest": task("t-unittest-157"),
+                "foobar/xpcshell-gpu": task("t-unittest-314-gpu"),
+                "foobar/xpcshell": task("t-unittest-314"),
+            },
+            {
+                "try_task_config": {
+                    "worker-types": ["t-unittest-314", "t-unittest-314-gpu"]
+                }
+            },
+            [
+                "foobar/xpcshell-1",
+                "foobar/xpcshell-gpu",
+                "foobar/xpcshell",
+            ],
+            id="multiple workers worker",
+        ),
+        pytest.param(
+            {
+                "foobar/xpcshell-1": task("t-unittest-314"),
+                "foobar/mochitest": task("t-unittest-157"),
+                "foobar/xpcshell-gpu": task("t-unittest-314-gpu"),
+                "foobar/xpcshell": task("t-unittest-314"),
+            },
+            {"try_task_config": {"worker-types": ["t-unittest-157"]}},
+            [
+                "foobar/mochitest",
+            ],
+            id="single task",
+        ),
+        pytest.param(
+            {
+                "foobar/xpcshell-1": task("t-unittest-314"),
+                "foobar/mochitest": task("t-unittest-157"),
+                "foobar/xpcshell-gpu": task("t-unittest-314-gpu"),
+                "foobar/xpcshell": task("t-unittest-314"),
+            },
+            {"try_task_config": {"worker-types": []}},
+            [
+                "foobar/xpcshell-1",
+                "foobar/mochitest",
+                "foobar/xpcshell-gpu",
+                "foobar/xpcshell",
+            ],
+            id="no worker",
+        ),
+        pytest.param(
+            {
+                "foobar/xpcshell-1": task("t-unittest-314"),
+                "foobar/mochitest": task("t-unittest-157"),
+                "foobar/xpcshell-gpu": task("t-unittest-314-gpu"),
+                "foobar/xpcshell": task("t-unittest-314"),
+            },
+            {"try_task_config": {"worker-types": ["fake-worker"]}},
+            [],
+            id="invalid worker",
+        ),
+    ),
+)
+def test_filter_tasks_by_worker_type(patch_resolver, tasks, params, expected):
+    assert list(filter_tasks_by_worker_type(tasks, params)) == expected
 
 
 def test_filter_tasks_by_paths(patch_resolver):