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

Code before:

            def run(self):
                super(TestWrapperTask, self).run()
                yield e2, e1

        wt = TestWrapperTask()

        with Worker(scheduler=self.sch, worker_id='X', keep_alive=True, wait_interval=0.1, wait_jitter=0.05) as w1:
            with Worker(scheduler=self.sch, worker_id='Y', keep_alive=True, wait_interval=0.1, wait_jitter=0.05) as w2:
                self.assertTrue(w1.add(wt))
                self.assertTrue(w2.add(s1))

                self.assertTrue(w2.run())
                self.assertFalse(w1.run())

                self.assertEqual([wt.task_id], list(self.sch.task_list('PENDING', 'UPSTREAM_DISABLED').keys()))

                self.assertEqual(sorted([e1.task_id, e2.task_id]), sorted(self.sch.task_list('DISABLED', '').keys()))

                self.assertEqual(0, self.sch._state.get_task(wt.task_id).failures.num_failures())
                self.assertEqual(0, self.sch._state.get_task(s1.task_id).failures.num_failures())
                self.assertEqual(self.per_task_retry_count, self.sch._state.get_task(e2.task_id).failures.num_failures())
                self.assertEqual(self.default_retry_count, self.sch._state.get_task(e1.task_id).failures.num_failures())


class SleepyLocalTarget(luigi.LocalTarget):
    time_to_sleep = 0.5

    def exists(self):
        time.sleep(self.time_to_sleep)
        return super(luigi.LocalTarget, self).exists()

    def touch(self):
        if os.path.exists(self.path):
            raise Exception("Path {} already exists".format(self.path))
        else:
            self.open("w").close()


class Child(luigi.Task):
    time_to_sleep = 1
    output_dir = luigi.Parameter(default="")

    def output(self):
        return SleepyLocalTarget("{}/Child".format(self.output_dir))

    def run(self):
        time.sleep(self.time_to_sleep)
        self.output().touch()



Review comment (performance):
Can we reduce the complexity of the testcase you're building? I'm not sure I fully understand what you're solving and that maybe could be easier if the test case was simpler (if possible).

Code after:

        class TestErrorTask1(DummyErrorTask):
            pass

        e1 = TestErrorTask1()

        class TestErrorTask2(DummyErrorTask):
            retry_count = self.per_task_retry_count

        e2 = TestErrorTask2()

        class TestSuccessTask1(DummyTask):
            pass

        s1 = TestSuccessTask1()

        class TestWrapperTask(DummyTask):
            def requires(self):
                return [s1]

            def run(self):
                super(TestWrapperTask, self).run()
                yield e2, e1

        wt = TestWrapperTask()

        with Worker(scheduler=self.sch, worker_id='X', keep_alive=True, wait_interval=0.1, wait_jitter=0.05) as w1:
            with Worker(scheduler=self.sch, worker_id='Y', keep_alive=True, wait_interval=0.1, wait_jitter=0.05) as w2:
                self.assertTrue(w1.add(wt))
                self.assertTrue(w2.add(s1))

                self.assertTrue(w2.run())
                self.assertFalse(w1.run())

                self.assertEqual([wt.task_id], list(self.sch.task_list('PENDING', 'UPSTREAM_DISABLED').keys()))

                self.assertEqual(sorted([e1.task_id, e2.task_id]), sorted(self.sch.task_list('DISABLED', '').keys()))

                self.assertEqual(0, self.sch._state.get_task(wt.task_id).failures.num_failures())
                self.assertEqual(0, self.sch._state.get_task(s1.task_id).failures.num_failures())
                self.assertEqual(self.per_task_retry_count, self.sch._state.get_task(e2.task_id).failures.num_failures())
                self.assertEqual(self.default_retry_count, self.sch._state.get_task(e1.task_id).failures.num_failures())