1# Copyright 2017 The Abseil Authors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""Helper script used by app_test.py."""
16
17import os
18import sys
19
20try:
21  import faulthandler  # pylint: disable=g-import-not-at-top
22except ImportError:
23  faulthandler = None
24
25from absl import app  # pylint: disable=g-import-not-at-top
26from absl import flags
27
28FLAGS = flags.FLAGS
29flags.DEFINE_boolean('faulthandler_sigsegv', False, 'raise SIGSEGV')
30flags.DEFINE_boolean('raise_exception', False, 'Raise MyException from main.')
31flags.DEFINE_boolean(
32    'raise_usage_error', False, 'Raise app.UsageError from main.')
33flags.DEFINE_integer(
34    'usage_error_exitcode', None, 'The exitcode if app.UsageError if raised.')
35flags.DEFINE_string(
36    'str_flag_with_unicode_args', u'thumb:\U0001F44D', u'smile:\U0001F604')
37flags.DEFINE_boolean('print_init_callbacks', False,
38                     'print init callbacks and exit')
39
40
41class MyException(Exception):
42  pass
43
44
45class MyExceptionHandler(app.ExceptionHandler):
46
47  def __init__(self, message):
48    self.message = message
49
50  def handle(self, exc):
51    sys.stdout.write('MyExceptionHandler: {}\n'.format(self.message))
52
53
54def real_main(argv):
55  """The main function."""
56  if os.environ.get('APP_TEST_PRINT_ARGV', False):
57    sys.stdout.write('argv: {}\n'.format(' '.join(argv)))
58
59  if FLAGS.raise_exception:
60    raise MyException
61
62  if FLAGS.raise_usage_error:
63    if FLAGS.usage_error_exitcode is not None:
64      raise app.UsageError('Error!', FLAGS.usage_error_exitcode)
65    else:
66      raise app.UsageError('Error!')
67
68  if FLAGS.faulthandler_sigsegv:
69    faulthandler._sigsegv()  # pylint: disable=protected-access
70    sys.exit(1)  # Should not reach here.
71
72  if FLAGS.print_init_callbacks:
73    app.call_after_init(lambda: _callback_results.append('during real_main'))
74    for value in _callback_results:
75      print('callback: {}'.format(value))
76    sys.exit(0)
77
78  # Ensure that we have a random C++ flag in flags.FLAGS; this shows
79  # us that app.run() did the right thing in conjunction with C++ flags.
80  helper_type = os.environ['APP_TEST_HELPER_TYPE']
81  if helper_type == 'clif':
82    if 'heap_check_before_constructors' in flags.FLAGS:
83      print('PASS: C++ flag present and helper_type is {}'.format(helper_type))
84      sys.exit(0)
85    else:
86      print('FAILED: C++ flag absent but helper_type is {}'.format(helper_type))
87      sys.exit(1)
88  elif helper_type == 'pure_python':
89    if 'heap_check_before_constructors' in flags.FLAGS:
90      print('FAILED: C++ flag present but helper_type is pure_python')
91      sys.exit(1)
92    else:
93      print('PASS: C++ flag absent and helper_type is pure_python')
94      sys.exit(0)
95  else:
96    print('Unexpected helper_type "{}"'.format(helper_type))
97    sys.exit(1)
98
99
100def custom_main(argv):
101  print('Function called: custom_main.')
102  real_main(argv)
103
104
105def main(argv):
106  print('Function called: main.')
107  real_main(argv)
108
109
110flags_parser_argv_sentinel = object()
111
112
113def flags_parser_main(argv):
114  print('Function called: main_with_flags_parser.')
115  if argv is not flags_parser_argv_sentinel:
116    sys.exit(
117        'FAILED: main function should be called with the return value of '
118        'flags_parser, but found: {}'.format(argv))
119
120
121def flags_parser(argv):
122  print('Function called: flags_parser.')
123  if os.environ.get('APP_TEST_FLAGS_PARSER_PARSE_FLAGS', None):
124    FLAGS(argv)
125  return flags_parser_argv_sentinel
126
127
128# Holds results from callbacks triggered by `app.run_after_init`.
129_callback_results = []
130
131if __name__ == '__main__':
132  kwargs = {'main': main}
133  main_function_name = os.environ.get('APP_TEST_CUSTOM_MAIN_FUNC', None)
134  if main_function_name:
135    kwargs['main'] = globals()[main_function_name]
136  custom_argv = os.environ.get('APP_TEST_CUSTOM_ARGV', None)
137  if custom_argv:
138    kwargs['argv'] = custom_argv.split(' ')
139  if os.environ.get('APP_TEST_USE_CUSTOM_PARSER', None):
140    kwargs['flags_parser'] = flags_parser
141
142  app.call_after_init(lambda: _callback_results.append('before app.run'))
143  app.install_exception_handler(MyExceptionHandler('first'))
144  app.install_exception_handler(MyExceptionHandler('second'))
145  app.run(**kwargs)
146
147  sys.exit('This is not reachable.')
148