Fix unit test

Signed-off-by: Huy Do <[email protected]>

Author: huydhn

.ci/scripts/tests/test_gather_benchmark_configs.py

@@ -192,20 +192,28 @@ def test_set_output_no_github_env(self, mock_getenv, mock_file):
 
     def test_device_pools_contains_all_devices(self):
         expected_devices = [
-            "apple_iphone_15",
-            "apple_iphone_15+ios_18",
-            "samsung_galaxy_s22",
-            "samsung_galaxy_s24",
-            "google_pixel_8_pro",
+            "apple_iphone_15+public",
+            "apple_iphone_15+ios_18_public",
+            "samsung_galaxy_s22+public",
+            "samsung_galaxy_s24+ultra_private",
+            "google_pixel_8+pro_public",
         ]
         for device in expected_devices:
-            self.assertIn(device, self.gather_benchmark_configs.DEVICE_POOLS)
+            m = re.match(self.gather_benchmark_configs.DEVICE_POOLS_REGEX, device)
+
+            device_name = m.group("device_name")
+            variant = m.group("variant")
+
+            self.assertIn(device_name, self.gather_benchmark_configs.DEVICE_POOLS)
+            self.assertIn(
+                variant, self.gather_benchmark_configs.DEVICE_POOLS[device_name]
+            )
 
     def test_gather_benchmark_configs_cli(self):
         args = {
             "models": "mv2,dl3",
             "os": "ios",
-            "devices": "apple_iphone_15",
+            "devices": "apple_iphone_15+pro_private",
             "configs": None,
         }
 
@@ -223,11 +231,29 @@ def test_gather_benchmark_configs_cli(self):
         self.assertIn('"config": "xnnpack_q8"', result.stdout)
         self.assertIn('"config": "mps"', result.stdout)
 
-    def test_gather_benchmark_configs_cli_specified_configs(self):
+    def test_gather_benchmark_configs_cli_invalid_device(self):
         args = {
             "models": "mv2,dl3",
             "os": "ios",
             "devices": "apple_iphone_15",
+            "configs": None,
+        }
+
+        cmd = ["python", ".ci/scripts/gather_benchmark_configs.py"]
+        for key, value in args.items():
+            if value is not None:
+                cmd.append(f"--{key}")
+                cmd.append(value)
+
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        self.assertEqual(result.returncode, 0, f"Error: {result.stderr}")
+        self.assertIn('{"include": []}', result.stdout)
+
+    def test_gather_benchmark_configs_cli_specified_configs(self):
+        args = {
+            "models": "mv2,dl3",
+            "os": "ios",
+            "devices": "apple_iphone_15+private",
             "configs": "coreml_fp16,xnnpack_q8",
         }
 
@@ -249,7 +275,7 @@ def test_gather_benchmark_configs_cli_specified_configs_raise(self):
         args = {
             "models": "mv2,dl3",
             "os": "ios",
-            "devices": "apple_iphone_15",
+            "devices": "apple_iphone_15+public",
             "configs": "qnn_q8",
         }