Add test cases

Author: KKIEEK

tests/test_ray/test_spaces.py

@@ -1,10 +1,8 @@
+import pytest
 from ray import tune
 
-from mmtune.ray.spaces import BaseSpace, Choice, build_space
-
-
-def test_base_space():
-    assert hasattr(BaseSpace(), 'space')
+from mmtune.ray.spaces import Choice, GridSearch, SampleFrom, build_space
+from mmtune.utils import ImmutableContainer
 
 
 def test_build_space():
@@ -18,11 +16,58 @@ def test_build_space():
 
 def test_choice():
 
-    def objective(config):
-        assert config.get('test') in [1, 2, 3]
-        tune.report(result=config['test'])
+    def is_in(config):
+        assert config['test'] in [0, 1, 2]
+
+    choice = Choice(categories=[0, 1, 2])
+    tune.run(is_in, config=dict(test=choice.space))
+
+    # with alias
+    def is_immutable(config):
+        assert isinstance(config['test'], ImmutableContainer)
+        assert config['test'].data in [True, False]
+        assert config['test'].alias in ['T', 'F']
+
+    with pytest.raises(AssertionError):
+        choice = Choice(categories=[True, False], alias=['TF'])
+    choice = Choice(categories=[True, False], alias=['T', 'F'])
+    tune.run(is_immutable, config=dict(test=choice.space))
+
+
+def test_grid_search():
+
+    def is_in(config):
+        print(config)
+        assert config['test1'] in [0, 1, 2]
+        assert config['test2'] in [3, 4, 5]
+
+    grid1 = GridSearch(values=[0, 1, 2])
+    grid2 = GridSearch(values=[3, 4, 5])
+    tune.run(is_in, config=dict(test1=grid1.space, test2=grid2.space))
+
+    # with alias
+    def is_immutable(config):
+        for test in ['test1', 'test2']:
+            assert isinstance(config[test], ImmutableContainer)
+            assert config[test].data in [True, False]
+            assert config[test].alias in ['T', 'F']
+
+    with pytest.raises(AssertionError):
+        grid1 = GridSearch(values=[True, False], alias=['TF'])
+    grid1 = GridSearch(values=[True, False], alias=['T', 'F'])
+    grid2 = GridSearch(values=[False, True], alias=['F', 'T'])
+    tune.run(is_immutable, config=dict(test1=grid1.space, test2=grid2.space))
+
+
+def test_sample_from():
+
+    def is_eq(config):
+        assert config['test'] == config['base']**2
+
+    with pytest.raises(AssertionError):
+        sample_from = SampleFrom('wrong expression')
+    sample_from = SampleFrom('lambda spec: spec.config.base ** 2')
+    tune.run(is_eq, config=dict(base=10, test=sample_from.space))
 
-    tune.run(
-        objective,
-        config=dict(
-            test=Choice(categories=[1, 2, 3], use_container=False).space))
+    sample_from = SampleFrom(lambda spec: spec.config.base**2)
+    tune.run(is_eq, config=dict(base=10, test=sample_from.space))