wip

nose_launchable/client.py

@@ -1,11 +1,11 @@
 import os
 
 import requests
+import subprocess
 
 from nose_launchable.log import logger
 from nose_launchable.version import __version__
 
-
 class LaunchableClientFactory:
     BASE_URL_KEY = "LAUNCHABLE_BASE_URL"
     TOKEN_KEY = "LAUNCHABLE_TOKEN"
@@ -16,7 +16,7 @@ class LaunchableClientFactory:
     def prepare(cls):
         url, org, wp, token = cls._parse_options()
 
-        return LaunchableClient(url, org, wp, token, requests)
+        return LaunchableClient(url, org, wp, token, requests, subprocess)
 
     @classmethod
     def _parse_options(cls):
@@ -35,12 +35,13 @@ def _get_base_url(cls):
 class LaunchableClient:
     CLIENT_NAME = "nose-launchable"
 
-    def __init__(self, base_url, org_name, workspace_name, token, http):
+    def __init__(self, base_url, org_name, workspace_name, token, http, process):
         self.base_url = base_url
         self.org_name = org_name
         self.workspace_name = workspace_name
         self.token = token
         self.http = http
+        self.process = process
         self.build_number = None
         self.test_session_id = None
 
@@ -83,6 +84,27 @@ def infer(self, test):
 
         return response_body
 
+    def subset(self, test_names, target):
+        url = "/test_sessions/{}".format(self.test_session_id)
+
+        proc = self.process.run(
+            ['launchable', 'subset', '--session', url, '--target', target + '%', 'file'],
+            input="\n".join(test_names),
+            encoding='utf-8',
+            stdout=self.process.PIPE,
+            stderr=self.process.PIPE
+        )
+
+        if proc.returncode != 0:
+            raise RuntimeError("launchable subset command fails. stdout: {}, stderr: {}", proc.stdout, proc.stderr)
+
+        # launchable subset command returns a list of test names splitted by \n
+        order = proc.stdout.rstrip("\n").split("\n")
+
+        logger.debug("Subset test order: {}".format(order))
+
+        return order
+
     def upload_events(self, events):
         url = "{}/intake/organizations/{}/workspaces/{}/builds/{}/test_sessions/{}/events".format(
             self.base_url,

nose_launchable/manager.py

@@ -36,6 +36,24 @@ def dfs(suite):
     return {"type": TREE_TYPE, "root": node}
 
 
+# Parse tests and return a list of test names
+def get_test_names(test):
+    test_names = []
+
+    def dfs(suite):
+        logger.debug("Parsing a test tree: suite: {}".format(suite))
+        if _is_leaf(suite):
+            test_names.append(_get_test_name(suite))
+            return suite
+
+        # Access to _tests is through a generator, so iteration is not repeatable by default
+        suite._tests = [dfs(t) for t in suite]
+        return suite
+
+    dfs(test)
+    return test_names
+
+
 # Reorder tests based on the given order
 def reorder(suite, order):
     if len(order) == 0:
@@ -83,6 +101,45 @@ def _reorder(suite, order):
     return tree_nodes
 
 
+# Subset tests based on the given order
+def subset(test, order):
+    def dfs(suite):
+        logger.debug("Subsetting a test tree: suite: {}".format(suite))
+        if _is_leaf(suite):
+            name = _get_test_name(suite)
+
+            score, is_target = 0, name in order
+            if is_target:
+                score = order[name]
+
+            logger.debug("A leaf node: score: {}, is_target: {}, suite: {}".format(score, is_target, suite))
+            return score, is_target, suite
+
+        cases = []
+        score = 0
+
+        for t in suite:
+            s, is_target, c = dfs(t)
+
+            if not is_target:
+                continue
+
+            cases.append((s, c))
+            score += s
+
+        # The smaller the score is, the faster the test needs to be tested
+        suite._tests = [c for _, c in sorted(cases, key=lambda x: x[0])]
+
+        # Propagate a score to its parent by calculating its average
+        # Avoid ZeroDivisionError with max(len(cases), 1)
+        s, is_target = score / max(len(cases), 1), len(cases) != 0
+        logger.debug("A non-leaf node: score: {}, is_target: {}, suite: {}".format(s, is_target, suite))
+        return s, is_target, suite
+
+    dfs(test)
+    return test
+
+
 # Check if test's context is a leaf
 def _is_leaf(suite):
     if type(suite) is ContextSuite:

nose_launchable/plugin.py

@@ -12,7 +12,7 @@
 from nose_launchable.case_event import CaseEvent
 from nose_launchable.client import LaunchableClientFactory
 from nose_launchable.log import logger
-from nose_launchable.manager import parse_test, reorder
+from nose_launchable.manager import parse_test, get_test_names, reorder, subset
 from nose_launchable.protecter import protect
 from nose_launchable.uploader import UploaderFactory
 
@@ -34,18 +34,35 @@ def __init__(self):
 
     def options(self, parser, env):
         super(Launchable, self).options(parser, env=env)
-        parser.add_option("--launchable", action='store_true', dest="enabled", help="Enable Launchable API interaction")
+        parser.add_option("--launchable", action='store_true', dest="enabled", help="Enable Launchable reordering")
+        parser.add_option("--launchable-reorder", action='store_true', dest="reorder_enabled", help="Enable Launchable reordering")
+        parser.add_option("--launchable-subset", action='store_true', dest="subset_enabled", help="Enable Launchable subsetting")
         parser.add_option("--launchable-build-number", action='store', type='string', dest="build_number", help="CI/CD build number")
+        parser.add_option("--launchable-subset-target", action='store', type='string', dest="subset_target", help="Target percentage of subset")
 
     def configure(self, options, conf):
         super(Launchable, self).configure(options, conf)
-        self.enabled = options.enabled
+
+        self.reorder_enabled = options.enabled or options.reorder_enabled or False
+        self.subset_enabled = options.subset_enabled or False
         self.build_number = options.build_number or os.getenv(BUILD_NUMBER_KEY)
+        self.subset_target = options.subset_target
+
+        if not (self.reorder_enabled ^ self.subset_enabled):
+            self.enabled = False
+            logger.warning("Please specify either --launchable flag or --launchable-subset flag")
+            return
+
+        self.enabled = True
 
         if self.enabled and self.build_number is None:
             self.enabled = False
-            logger.warning("--launchable flag is specified but --launchable-build-number flag is missing. "
-                           "Please specify --launchable-build-number flag in order to enable nose-launchable plugin")
+            logger.warning("Please specify --launchable-build-number flag in order to enable nose-launchable plugin")
+
+        if self.subset_enabled and self.subset_target is None:
+            self.enabled = False
+            logger.warning("Please specify --launchable-subset-target flag to run subset")
+
 
     @protect
     def begin(self):
@@ -58,13 +75,12 @@ def begin(self):
 
     @protect
     def prepareTest(self, test):
-        t = parse_test(test)
-
         self._print("Getting optimized test execution order from Launchable...\n")
-        order = self._client.infer(t)
-        self._print("Received optimized test execution order from Launchable\n")
+        if self.reorder_enabled:
+            self._reorder(test)
 
-        reorder(test, order)
+        if self.subset_enabled:
+            self._subset(test)
 
         self._print("Test execution optimized by Launchable ")
         # A rocket emoji
@@ -112,6 +128,21 @@ def finalize(self, test):
 
         self._uploader.join()
 
+    def _reorder(self, test):
+        tree = parse_test(test)
+
+        order = self._client.reorder(tree)
+
+        reorder(test, order)
+
+    def _subset(self, test):
+        test_names = get_test_names(test)
+
+        order = self._client.subset(test_names, self.subset_target)
+
+        # Create a dictionary maps test_name to its index
+        subset(test, {o: i for i, o in enumerate(order)})
+
     def _startCapture(self):
         self._capture_stack.append((sys.stdout, sys.stderr))
         self._currentStdout = StringIO()

requirements.txt

@@ -1,3 +1,4 @@
+launchable>=1.1.3
 nose>=1.0.0
 boto3>=1.0.0
 requests>=2.0.0

tests/test_client.py

@@ -1,11 +1,12 @@
 import os
+import subprocess
 import unittest
 from unittest.mock import MagicMock
 
 import requests
 
-from nose_launchable.client import LaunchableClientFactory, LaunchableClient
 from nose_launchable.case_event import CaseEvent
+from nose_launchable.client import LaunchableClientFactory, LaunchableClient
 from nose_launchable.version import __version__
 
 
@@ -21,6 +22,7 @@ def test_prepare(self):
         self.assertEqual('wp_name', client.workspace_name)
         self.assertEqual('v1:org_name/wp_name:token', client.token)
         self.assertEqual(requests, client.http)
+        self.assertEqual(subprocess, client.process)
 
     def test_prepare_with_base_url(self):
         os.environ[LaunchableClientFactory.BASE_URL_KEY] = 'base_url'
@@ -32,6 +34,7 @@ def test_prepare_with_base_url(self):
         self.assertEqual('wp_name', client.workspace_name)
         self.assertEqual('v1:org_name/wp_name:token', client.token)
         self.assertEqual(requests, client.http)
+        self.assertEqual(subprocess, client.process)
 
 
 class TestLaunchableClient(unittest.TestCase):
@@ -41,7 +44,9 @@ def test_start(self):
         mock_response.json.return_value = {'id': 1}
         mock_requests.post.return_value = mock_response
 
-        client = LaunchableClient("base_url", "org_name", "wp_name", "token", mock_requests)
+        mock_subprocess = MagicMock(name="subprecess")
+
+        client = LaunchableClient("base_url", "org_name", "wp_name", "token", mock_requests, mock_subprocess)
         client.start("test_build_number")
 
         expected_url = "base_url/intake/organizations/org_name/workspaces/wp_name/builds/test_build_number/test_sessions"
@@ -65,7 +70,9 @@ def test_infer(self):
         mock_requests = MagicMock(name="requests")
         mock_requests.post.return_value = mock_response
 
-        client = LaunchableClient("base_url", "org_name", "wp_name", "token", mock_requests)
+        mock_subprocess = MagicMock(name="subprecess")
+
+        client = LaunchableClient("base_url", "org_name", "wp_name", "token", mock_requests, mock_subprocess)
         client.build_number = "test"
         client.test_session_id = 1
 
@@ -96,12 +103,55 @@ def test_infer(self):
         mock_response.raise_for_status.assert_called_once_with()
         mock_response.json.assert_called_once_with()
 
+    def test_subset_success(self):
+        mock_output = MagicMock(name="output")
+        mock_subprocess = MagicMock(name="subprecess")
+
+        mock_subprocess.run.return_value = mock_output
+        mock_subprocess.PIPE = "PIPE"
+        # Success
+        mock_output.returncode = 0
+        mock_output.stdout = "tests/test2.py\ntests/test1.py\n"
+
+        mock_requests = MagicMock(name="requests")
+
+        client = LaunchableClient("base_url", "org_name", "wp_name", "token", mock_requests, mock_subprocess)
+        client.test_session_id = 1
+
+        got = client.subset(["tests/test1.py", "tests/test2.py"], "10")
+
+        expected_command = ['launchable', 'subset', '--session', '/test_sessions/1', '--target', '10%', 'file']
+        expected_input = 'tests/test1.py\ntests/test2.py'
+
+        mock_subprocess.run.assert_called_once_with(expected_command, input=expected_input, encoding='utf-8', stdout='PIPE', stderr='PIPE')
+        self.assertEqual(['tests/test2.py', 'tests/test1.py'], got)
+
+    def test_subset_failure(self):
+        mock_output = MagicMock(name="output")
+        mock_subprocess = MagicMock(name="subprecess")
+
+        mock_subprocess.run.return_value = mock_output
+        mock_subprocess.PIPE = "PIPE"
+        # Fail
+        mock_output.returncode = 1
+        mock_output.error = "error"
+
+        mock_requests = MagicMock(name="requests")
+
+        client = LaunchableClient("base_url", "org_name", "wp_name", "token", mock_requests, mock_subprocess)
+        client.test_session_id = 1
+
+        with self.assertRaises(RuntimeError):
+            client.subset(["tests/test1.py", "tests/test2.py"], "10")
+
     def test_upload_events(self):
         mock_response = MagicMock(name="response")
         mock_requests = MagicMock(name="requests")
         mock_requests.post.return_value = mock_response
 
-        client = LaunchableClient("base_url", "org_name", "wp_name", "token", mock_requests)
+        mock_subprocess = MagicMock(name="subprecess")
+
+        client = LaunchableClient("base_url", "org_name", "wp_name", "token", mock_requests, mock_subprocess)
 
         events = [
             CaseEvent("test1", 0.1, CaseEvent.TEST_PASSED, "stdout1", "stderr1"),