Repository: spotify/luigi
File: test/task_running_resources_test.py (Python)

Code before:
            self.assertGreater(duration, min_duration)
            if max_duration > 0:
                self.assertLess(duration, max_duration)

    def test_remote_scheduler_serial(self):
        # serial test
        # run two tasks that do not reduce the "foo" resource
        # as the total foo resource (3) is smaller than the requirement of two tasks (4),
        # the scheduler is forced to run them serially which takes longer than 4 seconds
        task_a = ResourceTestTask(param="a")
        task_b = ResourceTestTask(param="b")

        with self.assert_duration(min_duration=4):
            luigi.build([task_a, task_b], self.factory, workers=2,
                        scheduler_port=self.get_http_port())

    def test_remote_scheduler_parallel(self):
        # parallel test
        # run two tasks and the first one lowers its requirement on the "foo" resource, so that
        # the total "foo" resource (3) is sufficient to run both tasks in parallel shortly after
        # the first task started, so the entire process should not exceed 4 seconds
        task_c = ResourceTestTask(param="c", reduce_foo=True)
        task_d = ResourceTestTask(param="d")

        with self.assert_duration(max_duration=4):
            luigi.build([task_c, task_d], self.factory, workers=2,
                        scheduler_port=self.get_http_port())

    def test_local_scheduler(self):
        # trivial local scheduler test
        # as each worker process would communicate with its own scheduler, concurrency tests with
        # multiple processes like the ones above don't make sense here, but one can at least test
        # the running_task_resources setter and getter
        sch = luigi.scheduler.Scheduler(resources={"foo": 2})

        with luigi.worker.Worker(scheduler=sch) as w:

            task = ResourceTestTask(param="x", reduce_foo=True)

            w.add(task)
            w.run()

            self.assertEqual(sch.get_running_task_resources(task.task_id)["resources"]["foo"], 1)


Review comment (suggestion):
I wonder if you get strange behavior because of the simplistic and unintuitive way `luigi.build` is implemented. I think it just runs those two guys serially (it won't attempt `task_d` before `task_c` is complete. But you might just want to test to build with a single WrapperTask that depends on both of those tasks.

Code after:
    def get_app(self):
        return luigi.server.app(luigi.scheduler.Scheduler())

    def setUp(self):
        super(ConcurrentRunningResourcesTest, self).setUp()

        # run the luigi server in a new process and wait for its startup
        self._process = multiprocessing.Process(target=luigi.server.run)
        self._process.start()
        time.sleep(0.5)

        # configure the rpc scheduler, update the foo resource
        self.sch = luigi.rpc.RemoteScheduler()
        self.sch.update_resource("foo", 3)

    def tearDown(self):
        super(ConcurrentRunningResourcesTest, self).tearDown()

        # graceful server shutdown
        self._process.terminate()
        self._process.join(timeout=1)
        if self._process.is_alive():
            os.kill(self._process.pid, signal.SIGKILL)

    @contextmanager
    def worker(self, scheduler=None, processes=2):
        with luigi.worker.Worker(scheduler=scheduler or self.sch, worker_processes=processes) as w:
            w._config.wait_interval = 0.2
            w._config.check_unfulfilled_deps = False
            yield w

    @contextmanager
    def assert_duration(self, min_duration=0, max_duration=-1):
        t0 = time.time()
        try:
            yield
        finally:
            duration = time.time() - t0
            self.assertGreater(duration, min_duration)
            if max_duration > 0:
                self.assertLess(duration, max_duration)

    def test_tasks_serial(self):
        # serial test
        # run two tasks that do not reduce the "foo" resource
        # as the total foo resource (3) is smaller than the requirement of two tasks (4),
        # the scheduler is forced to run them serially which takes longer than 4 seconds
        with self.worker() as w:
            w.add(ResourceWrapperTask(reduce_foo=False))

            with self.assert_duration(min_duration=4):