1# -*- coding=utf-8 -*-
2# Copyright 2017 The Abseil Authors.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Tests for absl.command_name."""
17
18import ctypes
19import errno
20import os
21import unittest
22from unittest import mock
23
24from absl import command_name
25from absl.testing import absltest
26
27
28def _get_kernel_process_name():
29  """Returns the Kernel's name for our process or an empty string."""
30  try:
31    with open('/proc/self/status', 'rt') as status_file:
32      for line in status_file:
33        if line.startswith('Name:'):
34          return line.split(':', 2)[1].strip().encode('ascii', 'replace')
35      return b''
36  except IOError:
37    return b''
38
39
40def _is_prctl_syscall_available():
41  try:
42    libc = ctypes.CDLL('libc.so.6', use_errno=True)
43  except OSError:
44    return False
45  zero = ctypes.c_ulong(0)
46  try:
47    status = libc.prctl(zero, zero, zero, zero, zero)
48  except AttributeError:
49    return False
50  if status < 0 and errno.ENOSYS == ctypes.get_errno():
51    return False
52  return True
53
54
55@unittest.skipIf(not _get_kernel_process_name(),
56                 '_get_kernel_process_name() fails.')
57class CommandNameTest(absltest.TestCase):
58
59  def assertProcessNameSimilarTo(self, new_name):
60    if not isinstance(new_name, bytes):
61      new_name = new_name.encode('ascii', 'replace')
62    actual_name = _get_kernel_process_name()
63    self.assertTrue(actual_name)
64    self.assertTrue(new_name.startswith(actual_name),
65                    msg='set {!r} vs found {!r}'.format(new_name, actual_name))
66
67  @unittest.skipIf(not os.access('/proc/self/comm', os.W_OK),
68                   '/proc/self/comm is not writeable.')
69  def test_set_kernel_process_name(self):
70    new_name = u'ProcessNam0123456789abcdefghijklmnöp'
71    command_name.set_kernel_process_name(new_name)
72    self.assertProcessNameSimilarTo(new_name)
73
74  @unittest.skipIf(not _is_prctl_syscall_available(),
75                   'prctl() system call missing from libc.so.6.')
76  def test_set_kernel_process_name_no_proc_file(self):
77    new_name = b'NoProcFile0123456789abcdefghijklmnop'
78    mock_open = mock.mock_open()
79    with mock.patch.object(command_name, 'open', mock_open, create=True):
80      mock_open.side_effect = IOError('mock open that raises.')
81      command_name.set_kernel_process_name(new_name)
82    mock_open.assert_called_with('/proc/self/comm', mock.ANY)
83    self.assertProcessNameSimilarTo(new_name)
84
85  def test_set_kernel_process_name_failure(self):
86    starting_name = _get_kernel_process_name()
87    new_name = b'NameTest'
88    mock_open = mock.mock_open()
89    mock_ctypes_cdll = mock.patch('ctypes.CDLL')
90    with mock.patch.object(command_name, 'open', mock_open, create=True):
91      with mock.patch('ctypes.CDLL') as mock_ctypes_cdll:
92        mock_open.side_effect = IOError('mock open that raises.')
93        mock_libc = mock.Mock(['prctl'])
94        mock_ctypes_cdll.return_value = mock_libc
95        command_name.set_kernel_process_name(new_name)
96    mock_open.assert_called_with('/proc/self/comm', mock.ANY)
97    self.assertEqual(1, mock_libc.prctl.call_count)
98    self.assertEqual(starting_name, _get_kernel_process_name())  # No change.
99
100  def test_make_process_name_useful(self):
101    test_name = 'hello.from.test'
102    with mock.patch('sys.argv', [test_name]):
103      command_name.make_process_name_useful()
104    self.assertProcessNameSimilarTo(test_name)
105
106
107if __name__ == '__main__':
108  absltest.main()
109