Fix Optuna test

kubeflow · Nov 2, 2023 · 27c2ddf · 27c2ddf
1 parent d35b7bc
commit 27c2ddf
Showing 1 changed file with 102 additions and 6 deletions.
diff --git a/test/unit/v1beta1/suggestion/test_optuna_service.py b/test/unit/v1beta1/suggestion/test_optuna_service.py
@@ -244,39 +244,114 @@ def test_get_suggestion(self, algorithm_name, algorithm_settings):
                 "cmaes",
                 {"restart_strategy": "ipop", "sigma": "0.1", "random_state": "10"},
                 20,
-                [],
+                [
+                    {
+                        "name": "param-1",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="5", min="1", list=[]
+                        ),
+                    },
+                    {
+                        "name": "param-2",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="10", min="9", list=[]
+                        ),
+                    },
+                ],
                 grpc.StatusCode.OK,
             ],
             # [CMAES] Invalid parameter name
             [
                 "cmaes",
                 {"invalid": "invalid", "sigma": "0.1"},
                 100,
-                [],
+                [
+                    {
+                        "name": "param-1",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="5", min="1", list=[]
+                        ),
+                    },
+                    {
+                        "name": "param-2",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="10", min="9", list=[]
+                        ),
+                    },
+                ],
                 grpc.StatusCode.INVALID_ARGUMENT,
             ],
             # [CMAES] Invalid restart_strategy
             [
                 "cmaes",
                 {"restart_strategy": "invalid", "sigma": "0.1"},
                 15,
-                [],
+                [
+                    {
+                        "name": "param-1",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="5", min="1", list=[]
+                        ),
+                    },
+                    {
+                        "name": "param-2",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="10", min="9", list=[]
+                        ),
+                    },
+                ],
                 grpc.StatusCode.INVALID_ARGUMENT,
             ],
             # [CMAES] Invalid sigma
             [
                 "cmaes",
                 {"restart_strategy": "None", "sigma": "-10"},
                 55,
-                [],
+                [
+                    {
+                        "name": "param-1",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="5", min="1", list=[]
+                        ),
+                    },
+                    {
+                        "name": "param-2",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="10", min="9", list=[]
+                        ),
+                    },
+                ],
                 grpc.StatusCode.INVALID_ARGUMENT,
             ],
             # [CMAES] Invalid random_state
             [
                 "cmaes",
                 {"sigma": "0.2", "random_state": "-20"},
                 25,
-                [],
+                [
+                    {
+                        "name": "param-1",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="5", min="1", list=[]
+                        ),
+                    },
+                    {
+                        "name": "param-2",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="10", min="9", list=[]
+                        ),
+                    },
+                ],
                 grpc.StatusCode.INVALID_ARGUMENT,
             ],
             # [CMAES] Invalid number of parameters
@@ -296,7 +371,28 @@ def test_get_suggestion(self, algorithm_name, algorithm_settings):
                 grpc.StatusCode.INVALID_ARGUMENT,
             ],
             # [RANDOM] Valid Case
-            ["random", {"random_state": "10"}, 23, [], grpc.StatusCode.OK],
+            [
+                "random",
+                {"random_state": "10"},
+                23,
+                [
+                    {
+                        "name": "param-1",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="5", min="1", list=[]
+                        ),
+                    },
+                    {
+                        "name": "param-2",
+                        "type": api_pb2.INT,
+                        "feasible_space": api_pb2.FeasibleSpace(
+                            max="10", min="9", list=[]
+                        ),
+                    },
+                ],
+                grpc.StatusCode.OK,
+            ],
             # [RANDOM] Invalid parameter name
             [
                 "random",