forked from spotify/luigi
/
cmdline_test.py
139 lines (104 loc) · 4.51 KB
/
cmdline_test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
# Copyright (c) 2012 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
import ConfigParser
import logging
import luigi
from luigi.mock import MockFile
import mock
import unittest
import warnings
class SomeTask(luigi.Task):
n = luigi.IntParameter()
def output(self):
return File('/tmp/test_%d' % self.n)
def run(self):
f = self.output().open('w')
f.write('done')
f.close()
class AmbiguousClass(luigi.Task):
pass
class AmbiguousClass(luigi.Task):
pass
class NonAmbiguousClass(luigi.ExternalTask):
pass
class NonAmbiguousClass(luigi.Task):
def run(self):
NonAmbiguousClass.has_run = True
class DontRegisterThisOne(luigi.Task):
register_cls = False
class TaskWithSameName(luigi.Task):
register_cls = False
def run(self):
self.x = 42
class TaskWithSameName(luigi.Task):
# there should be no ambiguity
def run(self):
self.x = 43
class CmdlineTest(unittest.TestCase):
def setUp(self):
global File
File = MockFile
MockFile._file_contents.clear()
def test_expose_deprecated(self):
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
luigi.expose(SomeTask)
self.assertEqual(w[-1].category, DeprecationWarning)
@mock.patch("logging.getLogger")
def test_cmdline_main_task_cls(self, logger):
luigi.run(['--local-scheduler', '--n', '100'], main_task_cls=SomeTask)
self.assertEqual(MockFile._file_contents, {'/tmp/test_100': 'done'})
@mock.patch("logging.getLogger")
def test_cmdline_other_task(self, logger):
luigi.run(['--local-scheduler', 'SomeTask', '--n', '1000'])
self.assertEqual(MockFile._file_contents, {'/tmp/test_1000': 'done'})
@mock.patch("logging.getLogger")
def test_cmdline_ambiguous_class(self, logger):
self.assertRaises(Exception, luigi.run, ['--local-scheduler', 'AmbiguousClass'])
@mock.patch("logging.getLogger")
@mock.patch("warnings.warn")
def test_cmdline_non_ambiguous_class(self, warn, logger):
luigi.run(['--local-scheduler', 'NonAmbiguousClass'])
self.assertTrue(NonAmbiguousClass.has_run)
@mock.patch("logging.getLogger")
@mock.patch("logging.StreamHandler")
def test_setup_interface_logging(self, handler, logger):
handler.return_value = mock.Mock(name="stream_handler")
with mock.patch("luigi.interface.setup_interface_logging.has_run", new=False):
luigi.interface.setup_interface_logging()
self.assertEqual([mock.call(handler.return_value)], logger.return_value.addHandler.call_args_list)
with mock.patch("luigi.interface.setup_interface_logging.has_run", new=False):
self.assertRaises(ConfigParser.NoSectionError, luigi.interface.setup_interface_logging, '/blah')
@mock.patch("warnings.warn")
@mock.patch("luigi.interface.setup_interface_logging")
def test_cmdline_logger(self, setup_mock, warn):
luigi.run(['Task', '--local-scheduler'])
self.assertEqual([mock.call(None)], setup_mock.call_args_list)
with mock.patch("luigi.configuration.get_config") as getconf:
getconf.return_value.get.return_value = None
getconf.return_value.get_boolean.return_value = True
luigi.interface.setup_interface_logging.call_args_list = []
luigi.run(['Task', '--local-scheduler'])
self.assertEqual([], setup_mock.call_args_list)
@mock.patch('argparse.ArgumentParser.print_usage')
def test_non_existent_class(self, print_usage):
self.assertRaises(SystemExit, luigi.run, ['--local-scheduler', 'XYZ'])
@mock.patch('argparse.ArgumentParser.print_usage')
def test_not_registered_class(self, print_usage):
self.assertRaises(SystemExit, luigi.run, ['--local-scheduler', 'DontRegisterThisOne'])
def test_unregistered_class(self):
luigi.run(['--local-scheduler', 'TaskWithSameName'])
self.assertEqual(TaskWithSameName().x, 43)
if __name__ == '__main__':
unittest.main()