1# Copyright 2020 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://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, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Tests for env_setup.environment.
15
16This tests the error-checking, context manager, and written environment scripts
17of the Environment class.
18
19Tests that end in "_ctx" modify the environment and validate it in-process.
20
21Tests that end in "_written" write the environment to a file intended to be
22evaluated by the shell, then launches the shell and then saves the environment.
23This environment is then validated in the test process.
24"""
25
26import logging
27import os
28import subprocess
29import tempfile
30import unittest
31
32import six
33
34from pw_env_setup import environment
35
36# pylint: disable=super-with-arguments
37
38
39class WrittenEnvFailure(Exception):
40    pass
41
42
43def _evaluate_env_in_shell(env):
44    """Write env to a file then evaluate and save the resulting environment.
45
46    Write env to a file, then launch a shell command that sources that file
47    and dumps the environment to stdout. Parse that output into a dict and
48    return it.
49
50    Args:
51      env(environment.Environment): environment to write out
52
53    Returns dictionary of resulting environment.
54    """
55
56    # Write env sourcing script to file.
57    with tempfile.NamedTemporaryFile(
58            prefix='pw-test-written-env-',
59            suffix='.bat' if os.name == 'nt' else '.sh',
60            delete=False,
61            mode='w+') as temp:
62        env.write(temp)
63        temp_name = temp.name
64
65    # Evaluate env sourcing script and capture output of 'env'.
66    if os.name == 'nt':
67        # On Windows you just run batch files and they modify your
68        # environment, no need to call 'source' or '.'.
69        cmd = '{} && set'.format(temp_name)
70    else:
71        # Using '.' instead of 'source' because 'source' is not POSIX.
72        cmd = '. {} && env'.format(temp_name)
73
74    res = subprocess.run(cmd, capture_output=True, shell=True)
75    if res.returncode:
76        raise WrittenEnvFailure(res.stderr)
77
78    # Parse environment from stdout of subprocess.
79    env_ret = {}
80    for line in res.stdout.splitlines():
81        line = line.decode()
82
83        # Some people inexplicably have newlines in some of their
84        # environment variables. This module does not allow that so we can
85        # ignore any such extra lines.
86        if '=' not in line:
87            continue
88
89        var, value = line.split('=', 1)
90        env_ret[var] = value
91
92    return env_ret
93
94
95# pylint: disable=too-many-public-methods
96class EnvironmentTest(unittest.TestCase):
97    """Tests for env_setup.environment."""
98    def setUp(self):
99        self.env = environment.Environment()
100
101        # Name of a variable that is already set when the test starts.
102        self.var_already_set = self.env.normalize_key('var_already_set')
103        os.environ[self.var_already_set] = 'orig value'
104        self.assertIn(self.var_already_set, os.environ)
105
106        # Name of a variable that is not set when the test starts.
107        self.var_not_set = self.env.normalize_key('var_not_set')
108        if self.var_not_set in os.environ:
109            del os.environ[self.var_not_set]
110        self.assertNotIn(self.var_not_set, os.environ)
111
112        self.orig_env = os.environ.copy()
113
114    def tearDown(self):
115        self.assertEqual(os.environ, self.orig_env)
116
117    def test_set_notpresent_ctx(self):
118        self.env.set(self.var_not_set, '1')
119        with self.env(export=False) as env:
120            self.assertIn(self.var_not_set, env)
121            self.assertEqual(env[self.var_not_set], '1')
122
123    def test_set_notpresent_written(self):
124        self.env.set(self.var_not_set, '1')
125        env = _evaluate_env_in_shell(self.env)
126        self.assertIn(self.var_not_set, env)
127        self.assertEqual(env[self.var_not_set], '1')
128
129    def test_set_present_ctx(self):
130        self.env.set(self.var_already_set, '1')
131        with self.env(export=False) as env:
132            self.assertIn(self.var_already_set, env)
133            self.assertEqual(env[self.var_already_set], '1')
134
135    def test_set_present_written(self):
136        self.env.set(self.var_already_set, '1')
137        env = _evaluate_env_in_shell(self.env)
138        self.assertIn(self.var_already_set, env)
139        self.assertEqual(env[self.var_already_set], '1')
140
141    def test_clear_notpresent_ctx(self):
142        self.env.clear(self.var_not_set)
143        with self.env(export=False) as env:
144            self.assertNotIn(self.var_not_set, env)
145
146    def test_clear_notpresent_written(self):
147        self.env.clear(self.var_not_set)
148        env = _evaluate_env_in_shell(self.env)
149        self.assertNotIn(self.var_not_set, env)
150
151    def test_clear_present_ctx(self):
152        self.env.clear(self.var_already_set)
153        with self.env(export=False) as env:
154            self.assertNotIn(self.var_already_set, env)
155
156    def test_clear_present_written(self):
157        self.env.clear(self.var_already_set)
158        env = _evaluate_env_in_shell(self.env)
159        self.assertNotIn(self.var_already_set, env)
160
161    def test_value_replacement(self):
162        self.env.set(self.var_not_set, '/foo/bar/baz')
163        self.env.add_replacement('FOOBAR', '/foo/bar')
164        buf = six.StringIO()
165        self.env.write(buf)
166        assert '/foo/bar' not in buf.getvalue()
167
168    def test_variable_replacement(self):
169        self.env.set('FOOBAR', '/foo/bar')
170        self.env.set(self.var_not_set, '/foo/bar/baz')
171        self.env.add_replacement('FOOBAR')
172        buf = six.StringIO()
173        self.env.write(buf)
174        print(buf.getvalue())
175        assert '/foo/bar/baz' not in buf.getvalue()
176
177    def test_nonglobal(self):
178        self.env.set(self.var_not_set, '1')
179        with self.env(export=False) as env:
180            self.assertIn(self.var_not_set, env)
181            self.assertNotIn(self.var_not_set, os.environ)
182
183    def test_global(self):
184        self.env.set(self.var_not_set, '1')
185        with self.env(export=True) as env:
186            self.assertIn(self.var_not_set, env)
187            self.assertIn(self.var_not_set, os.environ)
188
189    def test_set_badnametype(self):
190        with self.assertRaises(environment.BadNameType):
191            self.env.set(123, '123')
192
193    def test_set_badvaluetype(self):
194        with self.assertRaises(environment.BadValueType):
195            self.env.set('var', 123)
196
197    def test_prepend_badnametype(self):
198        with self.assertRaises(environment.BadNameType):
199            self.env.prepend(123, '123')
200
201    def test_prepend_badvaluetype(self):
202        with self.assertRaises(environment.BadValueType):
203            self.env.prepend('var', 123)
204
205    def test_append_badnametype(self):
206        with self.assertRaises(environment.BadNameType):
207            self.env.append(123, '123')
208
209    def test_append_badvaluetype(self):
210        with self.assertRaises(environment.BadValueType):
211            self.env.append('var', 123)
212
213    def test_set_badname_empty(self):
214        with self.assertRaises(environment.BadVariableName):
215            self.env.set('', '123')
216
217    def test_set_badname_digitstart(self):
218        with self.assertRaises(environment.BadVariableName):
219            self.env.set('123', '123')
220
221    def test_set_badname_equals(self):
222        with self.assertRaises(environment.BadVariableName):
223            self.env.set('foo=bar', '123')
224
225    def test_set_badname_period(self):
226        with self.assertRaises(environment.BadVariableName):
227            self.env.set('abc.def', '123')
228
229    def test_set_badname_hyphen(self):
230        with self.assertRaises(environment.BadVariableName):
231            self.env.set('abc-def', '123')
232
233    def test_set_empty_value(self):
234        with self.assertRaises(environment.EmptyValue):
235            self.env.set('var', '')
236
237    def test_set_newline_in_value(self):
238        with self.assertRaises(environment.NewlineInValue):
239            self.env.set('var', '123\n456')
240
241    def test_equal_sign_in_value(self):
242        with self.assertRaises(environment.BadVariableValue):
243            self.env.append(self.var_already_set, 'pa=th')
244
245
246class _PrependAppendEnvironmentTest(unittest.TestCase):
247    """Tests for env_setup.environment."""
248    def __init__(self, *args, **kwargs):
249        windows = kwargs.pop('windows', False)
250        pathsep = kwargs.pop('pathsep', os.pathsep)
251        allcaps = kwargs.pop('allcaps', False)
252        super(_PrependAppendEnvironmentTest, self).__init__(*args, **kwargs)
253        self.windows = windows
254        self.pathsep = pathsep
255        self.allcaps = allcaps
256
257        # If we're testing Windows behavior and actually running on Windows,
258        # actually launch a subprocess to evaluate the shell init script.
259        # Likewise if we're testing POSIX behavior and actually on a POSIX
260        # system. Tests can check self.run_shell_tests and exit without
261        # doing anything.
262        real_windows = (os.name == 'nt')
263        self.run_shell_tests = (self.windows == real_windows)
264
265    def setUp(self):
266        self.env = environment.Environment(windows=self.windows,
267                                           pathsep=self.pathsep,
268                                           allcaps=self.allcaps)
269
270        self.var_already_set = self.env.normalize_key('VAR_ALREADY_SET')
271        os.environ[self.var_already_set] = self.pathsep.join(
272            'one two three'.split())
273        self.assertIn(self.var_already_set, os.environ)
274
275        self.var_not_set = self.env.normalize_key('VAR_NOT_SET')
276        if self.var_not_set in os.environ:
277            del os.environ[self.var_not_set]
278        self.assertNotIn(self.var_not_set, os.environ)
279
280        self.orig_env = os.environ.copy()
281
282    def split(self, val):
283        return val.split(self.pathsep)
284
285    def tearDown(self):
286        self.assertEqual(os.environ, self.orig_env)
287
288
289# TODO(mohrr) remove disable=useless-object-inheritance once in Python 3.
290# pylint: disable=useless-object-inheritance
291class _AppendPrependTestMixin(object):
292    def test_prepend_present_ctx(self):
293        orig = os.environ[self.var_already_set]
294        self.env.prepend(self.var_already_set, 'path')
295        with self.env(export=False) as env:
296            self.assertEqual(env[self.var_already_set],
297                             self.pathsep.join(('path', orig)))
298
299    def test_prepend_present_written(self):
300        if not self.run_shell_tests:
301            return
302
303        orig = os.environ[self.var_already_set]
304        self.env.prepend(self.var_already_set, 'path')
305        env = _evaluate_env_in_shell(self.env)
306        self.assertEqual(env[self.var_already_set],
307                         self.pathsep.join(('path', orig)))
308
309    def test_prepend_notpresent_ctx(self):
310        self.env.prepend(self.var_not_set, 'path')
311        with self.env(export=False) as env:
312            self.assertEqual(env[self.var_not_set], 'path')
313
314    def test_prepend_notpresent_written(self):
315        if not self.run_shell_tests:
316            return
317
318        self.env.prepend(self.var_not_set, 'path')
319        env = _evaluate_env_in_shell(self.env)
320        self.assertEqual(env[self.var_not_set], 'path')
321
322    def test_append_present_ctx(self):
323        orig = os.environ[self.var_already_set]
324        self.env.append(self.var_already_set, 'path')
325        with self.env(export=False) as env:
326            self.assertEqual(env[self.var_already_set],
327                             self.pathsep.join((orig, 'path')))
328
329    def test_append_present_written(self):
330        if not self.run_shell_tests:
331            return
332
333        orig = os.environ[self.var_already_set]
334        self.env.append(self.var_already_set, 'path')
335        env = _evaluate_env_in_shell(self.env)
336        self.assertEqual(env[self.var_already_set],
337                         self.pathsep.join((orig, 'path')))
338
339    def test_append_notpresent_ctx(self):
340        self.env.append(self.var_not_set, 'path')
341        with self.env(export=False) as env:
342            self.assertEqual(env[self.var_not_set], 'path')
343
344    def test_append_notpresent_written(self):
345        if not self.run_shell_tests:
346            return
347
348        self.env.append(self.var_not_set, 'path')
349        env = _evaluate_env_in_shell(self.env)
350        self.assertEqual(env[self.var_not_set], 'path')
351
352    def test_remove_ctx(self):
353        self.env.set(self.var_not_set,
354                     self.pathsep.join(('path', 'one', 'path', 'two', 'path')))
355
356        self.env.append(self.var_not_set, 'path')
357        with self.env(export=False) as env:
358            self.assertEqual(env[self.var_not_set],
359                             self.pathsep.join(('one', 'two', 'path')))
360
361    def test_remove_written(self):
362        if not self.run_shell_tests:
363            return
364
365        if self.windows:  # TODO(pwbug/231) Re-enable for Windows.
366            return
367
368        self.env.set(self.var_not_set,
369                     self.pathsep.join(('path', 'one', 'path', 'two', 'path')))
370
371        self.env.append(self.var_not_set, 'path')
372        env = _evaluate_env_in_shell(self.env)
373        self.assertEqual(env[self.var_not_set],
374                         self.pathsep.join(('one', 'two', 'path')))
375
376    def test_remove_ctx_space(self):
377        self.env.set(self.var_not_set,
378                     self.pathsep.join(('pa th', 'one', 'pa th', 'two')))
379
380        self.env.append(self.var_not_set, 'pa th')
381        with self.env(export=False) as env:
382            self.assertEqual(env[self.var_not_set],
383                             self.pathsep.join(('one', 'two', 'pa th')))
384
385    def test_remove_written_space(self):
386        if not self.run_shell_tests:
387            return
388
389        if self.windows:  # TODO(pwbug/231) Re-enable for Windows.
390            return
391
392        self.env.set(self.var_not_set,
393                     self.pathsep.join(('pa th', 'one', 'pa th', 'two')))
394
395        self.env.append(self.var_not_set, 'pa th')
396        env = _evaluate_env_in_shell(self.env)
397        self.assertEqual(env[self.var_not_set],
398                         self.pathsep.join(('one', 'two', 'pa th')))
399
400    def test_remove_ctx_empty(self):
401        self.env.remove(self.var_not_set, 'path')
402        with self.env(export=False) as env:
403            self.assertNotIn(self.var_not_set, env)
404
405    def test_remove_written_empty(self):
406        if not self.run_shell_tests:
407            return
408
409        self.env.remove(self.var_not_set, 'path')
410        env = _evaluate_env_in_shell(self.env)
411        self.assertNotIn(self.var_not_set, env)
412
413
414class WindowsEnvironmentTest(_PrependAppendEnvironmentTest,
415                             _AppendPrependTestMixin):
416    def __init__(self, *args, **kwargs):
417        kwargs['pathsep'] = ';'
418        kwargs['windows'] = True
419        kwargs['allcaps'] = True
420        super(WindowsEnvironmentTest, self).__init__(*args, **kwargs)
421
422
423class PosixEnvironmentTest(_PrependAppendEnvironmentTest,
424                           _AppendPrependTestMixin):
425    def __init__(self, *args, **kwargs):
426        kwargs['pathsep'] = ':'
427        kwargs['windows'] = False
428        kwargs['allcaps'] = False
429        super(PosixEnvironmentTest, self).__init__(*args, **kwargs)
430        self.real_windows = (os.name == 'nt')
431
432
433class WindowsCaseInsensitiveTest(unittest.TestCase):
434    def test_lower_handling(self):
435        # This is only for testing case-handling on Windows. It doesn't make
436        # sense to run it on other systems.
437        if os.name != 'nt':
438            return
439
440        lower_var = 'lower_var'
441        upper_var = lower_var.upper()
442
443        if upper_var in os.environ:
444            del os.environ[upper_var]
445
446        self.assertNotIn(lower_var, os.environ)
447
448        env = environment.Environment()
449        env.append(lower_var, 'foo')
450        env.append(upper_var, 'bar')
451        with env(export=False) as env_:
452            self.assertNotIn(lower_var, env_)
453            self.assertIn(upper_var, env_)
454            self.assertEqual(env_[upper_var], 'foo;bar')
455
456
457if __name__ == '__main__':
458    import sys
459    logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)
460    unittest.main()
461