blob: 08a49d8b1b1c0933686d2e442116879c4ab3e175 [file] [log] [blame]
Allen Li464220f2017-09-12 17:14:22 -07001# Copyright 2017 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Kludges to support legacy Autotest code.
6
Allen Libca5d012017-10-20 13:44:36 -07007Autotest imports should be done by calling patch() first and then
8calling load(). patch() should only be called once from a script's main
9function.
Allen Li464220f2017-09-12 17:14:22 -070010"""
11
12from __future__ import absolute_import
13from __future__ import division
14from __future__ import print_function
15
16import imp
17import importlib
18import logging
19import os
20import site
21import sys
22
23import autotest_lib
24
25_AUTOTEST_DIR = autotest_lib.__path__[0]
26_SITEPKG_DIR = os.path.join(_AUTOTEST_DIR, 'site-packages')
27
28_setup_done = False
29
30logger = logging.getLogger(__name__)
31
32
Allen Libca5d012017-10-20 13:44:36 -070033def patch():
34 """Monkeypatch everything needed to get Autotest working.
35
36 This should be called before any calls to load(). Only the main
37 function in scripts should call this function.
38
39 Unlike monkeypatch() which is more low-level, this patches not just
40 imports, but also other things in Autotest what would generally be
41 needed to use it.
42 """
43 monkeypatch()
44
Allen Lif4b62ae2017-11-09 15:48:05 -080045 # Add chromite's third-party to the import path.
46 import chromite
47
Allen Libca5d012017-10-20 13:44:36 -070048 # Needed to set up Django environment variables.
49 load('frontend.setup_django_environment')
50
51 # Monkey patch package paths that Django uses to be absolute.
52 settings = load('frontend.settings')
53 settings.INSTALLED_APPS = (
54 'autotest_lib.frontend.afe',
55 'autotest_lib.frontend.tko',
56 'django.contrib.admin',
57 'django.contrib.auth',
58 'django.contrib.contenttypes',
59 'django.contrib.sessions',
60 'django.contrib.sites',
61 )
62
Allen Lif4b62ae2017-11-09 15:48:05 -080063 # drone_utility uses this
64 common = load('scheduler.common')
65 common.autotest_dir = _AUTOTEST_DIR
66
Allen Libca5d012017-10-20 13:44:36 -070067
Allen Li464220f2017-09-12 17:14:22 -070068def monkeypatch():
Allen Libca5d012017-10-20 13:44:36 -070069 """Monkeypatch Autotest imports.
Allen Li464220f2017-09-12 17:14:22 -070070
Allen Li2678d362017-10-11 16:59:07 -070071 This should be called before any calls to load(). Only the main
72 function in scripts should call this function.
Allen Li464220f2017-09-12 17:14:22 -070073
74 This should be called no more than once.
75
Allen Li2678d362017-10-11 16:59:07 -070076 This adds Autotest's site-packages to the import path and modifies
77 sys.meta_path so that all common.py imports are no-ops.
Allen Li464220f2017-09-12 17:14:22 -070078 """
79 global _setup_done
80 assert not _setup_done
81 site.addsitedir(_SITEPKG_DIR)
82 sys.meta_path.insert(0, _CommonRemovingFinder())
83 _setup_done = True
84
85
86class _CommonRemovingFinder(object):
87 """Python import finder that neuters Autotest's common.py
88
89 The common module is replaced with an empty module everywhere it is
90 imported. common.py should have only been imported for side
91 effects, so nothing should actually use the imported module.
92
93 See also https://www.python.org/dev/peps/pep-0302/
94 """
95
96 def find_module(self, fullname, path=None):
97 """Find module."""
98 del path # unused
99 if not self._is_autotest_common(fullname):
100 return None
101 logger.debug('Dummying out %s import', fullname)
102 return self
103
104 def _is_autotest_common(self, fullname):
105 return (fullname.partition('.')[0] == 'autotest_lib'
106 and fullname.rpartition('.')[-1] == 'common')
107
108 def load_module(self, fullname):
109 """Load module."""
Allen Libca5d012017-10-20 13:44:36 -0700110 if fullname in sys.modules: # pragma: no cover
Allen Li464220f2017-09-12 17:14:22 -0700111 return sys.modules[fullname]
112 mod = imp.new_module(fullname)
113 mod.__file__ = '<removed>'
114 mod.__loader__ = self
115 mod.__package__ = fullname.rpartition('.')[0]
116 sys.modules[fullname] = mod
117 return mod
118
119
120def load(name):
121 """Import module from autotest.
122
Allen Li2678d362017-10-11 16:59:07 -0700123 This enforces that monkeypatch() is called first.
Allen Licc153502017-09-15 16:08:25 -0700124
125 @param name: name of module as string, e.g., 'frontend.afe.models'
Allen Li464220f2017-09-12 17:14:22 -0700126 """
127 if not _setup_done:
128 raise ImportError('cannot load Autotest modules before monkeypatching')
129 relpath = name.lstrip('.')
130 return importlib.import_module('.%s' % relpath, package='autotest_lib')