Add type annotation to func args and return type

* test_experiment_builder_discrete_default_params * test_experiment_builder_continuous_default_params
dantp-ai · Jan 2, 2024 · 0805062 · 0805062
1 parent 2d3c637
commit 0805062
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/test/highlevel/test_experiment_builder.py b/test/highlevel/test_experiment_builder.py
@@ -8,6 +8,7 @@
     DDPGExperimentBuilder,
     DiscreteSACExperimentBuilder,
     DQNExperimentBuilder,
+    ExperimentBuilder,
     ExperimentConfig,
     IQNExperimentBuilder,
     PGExperimentBuilder,
@@ -33,7 +34,7 @@
         PGExperimentBuilder,
     ],
 )
-def test_experiment_builder_continuous_default_params(builder_cls):
+def test_experiment_builder_continuous_default_params(builder_cls: type[ExperimentBuilder]) -> None:
     env_factory = ContinuousTestEnvFactory()
     sampling_config = SamplingConfig(
         num_epochs=1,
@@ -62,7 +63,7 @@ def test_experiment_builder_continuous_default_params(builder_cls):
         IQNExperimentBuilder,
     ],
 )
-def test_experiment_builder_discrete_default_params(builder_cls):
+def test_experiment_builder_discrete_default_params(builder_cls: type[ExperimentBuilder]) -> None:
     env_factory = DiscreteTestEnvFactory()
     sampling_config = SamplingConfig(
         num_epochs=1,