(Testing) Test Helper to test components with a runner. (#511)

Summary:
Pull Request resolved: #511

Exposing test helper method for covering component behavior using Runner and scheduler that is has a parity of running cli.

More details: #370

Reviewed By: d4l3k

Differential Revision: D36948958

fbshipit-source-id: 03f4853aab5e66f1e38727b715f44aae0fa2d8c0

Author: kurman

torchx/components/component_test_base.py

@@ -15,14 +15,32 @@
 """
 
 import os
+import shutil
+import tempfile
+import time
 import unittest
 from types import ModuleType
+from typing import Any, Callable, Dict, Optional
+
+from torchx.runner import get_runner
+
+from torchx.specs import AppDef, AppStatus
 
 from torchx.specs.builders import _create_args_parser
 from torchx.specs.finder import get_component
 
 
 class ComponentTestCase(unittest.TestCase):
+    def setUp(self) -> None:
+        self.test_dir = tempfile.mkdtemp("torchx_component_test")
+
+        self.old_cwd = os.getcwd()
+        os.chdir(os.path.dirname(os.path.dirname(__file__)))
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.test_dir)
+        os.chdir(self.old_cwd)
+
     """
     ComponentTestCase is an extension of TestCase with helper methods for use
     with testing component definitions.
@@ -58,3 +76,43 @@ def validate(self, module: ModuleType, function_name: str) -> None:
         # this will raise an exception and the test will fail
         with self.assertRaises(SystemExit):
             _ = _create_args_parser(component_def.fn).parse_args(["--help"])
+
+    def run_component(
+        self,
+        component: Callable[..., AppDef],
+        args: Optional[Dict[str, Any]] = None,
+        scheduler_params: Optional[Dict[str, Any]] = None,
+        scheduler: str = "local_cwd",
+        interval: float = 0.1,
+        timeout: float = 1,
+    ) -> Optional[AppStatus]:
+        """
+        Helper function that hides complexity of setting up the runner and polling results.
+        Note: method is blocking until either scheduler exits or timeout is reached (for non-blocking schedulers).
+
+        Args:
+                components: component function, factory for AppDef
+                args: optional component factory arguments
+                scheduler_params: optional parameters for scheduler factory method
+                scheduler: scheduler name
+                interval: scheduler comppletion polling interval
+                timeout: max time for scheduler to complete
+
+        """
+
+        app_def = component(**args)
+        if scheduler_params:
+            runner = get_runner(name=None, component_defaults=None, **scheduler_params)
+        else:
+            runner = get_runner(name=None, component_defaults=None)
+
+        app_handle = runner.run(app_def, scheduler)
+
+        elapsed = 0
+        status = runner.status(app_handle)
+        while (status and not status.is_terminal) or elapsed < timeout:
+            time.sleep(interval)
+            elapsed += interval
+            status = runner.status(app_handle)
+
+        return status

torchx/components/test/utils_test.py

@@ -6,6 +6,7 @@
 
 import torchx.components.utils as utils
 from torchx.components.component_test_base import ComponentTestCase
+from torchx.specs import AppState
 
 
 class UtilsComponentTest(ComponentTestCase):
@@ -29,3 +30,18 @@ def test_copy(self) -> None:
 
     def test_booth(self) -> None:
         self.validate(utils, "booth")
+
+    def test_run_sh(self) -> None:
+        result = self.run_component(
+            utils.echo, {"msg": "from test"}, scheduler_params={"cache_size": 1}
+        )
+        self.assertIsNotNone(result)
+        self.assertEqual(result.state, AppState.SUCCEEDED)
+
+    def test_run_sh_scheduler_factory_failure(self) -> None:
+        self.assertRaises(
+            ValueError,
+            lambda: self.run_component(
+                utils.echo, {"msg": "from test"}, scheduler_params={"cache_size": -1}
+            ),
+        )