bpo-38234: Add tests for Python init path config (GH-16358)

3 files changed, 176 insertions, 24 deletions

diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
index 2a5ace0..7808a1a 100644
--- a/Lib/test/test_embed.py
+++ b/Lib/test/test_embed.py
@@ -470,7 +470,8 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
                 xoptions[opt] = True
         return xoptions
 
-    def _get_expected_config(self, env):
+    def _get_expected_config_impl(self):
+        env = remove_python_envvars()
         code = textwrap.dedent('''
             import json
             import sys
@@ -499,13 +500,19 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
         except json.JSONDecodeError:
             self.fail(f"fail to decode stdout: {stdout!r}")
 
+    def _get_expected_config(self):
+        cls = InitConfigTests
+        if cls.EXPECTED_CONFIG is None:
+            cls.EXPECTED_CONFIG = self._get_expected_config_impl()
+
+        # get a copy
+        return {key: dict(value)
+                for key, value in cls.EXPECTED_CONFIG.items()}
+
     def get_expected_config(self, expected_preconfig, expected, env, api,
                             modify_path_cb=None):
         cls = self.__class__
-        if cls.EXPECTED_CONFIG is None:
-            cls.EXPECTED_CONFIG = self._get_expected_config(env)
-        configs = {key: dict(value)
-                   for key, value in self.EXPECTED_CONFIG.items()}
+        configs = self._get_expected_config()
 
         pre_config = configs['pre_config']
         for key, value in expected_preconfig.items():
@@ -553,9 +560,10 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
             if value is self.GET_DEFAULT_CONFIG:
                 expected[key] = config[key]
 
-        prepend_path = expected['pythonpath_env']
-        if prepend_path is not None:
-            expected['module_search_paths'] = [prepend_path, *expected['module_search_paths']]
+        pythonpath_env = expected['pythonpath_env']
+        if pythonpath_env is not None:
+            paths = pythonpath_env.split(os.path.pathsep)
+            expected['module_search_paths'] = [*paths, *expected['module_search_paths']]
         if modify_path_cb is not None:
             expected['module_search_paths'] = expected['module_search_paths'].copy()
             modify_path_cb(expected['module_search_paths'])
@@ -604,8 +612,11 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
 
     def check_all_configs(self, testname, expected_config=None,
                      expected_preconfig=None, modify_path_cb=None, stderr=None,
-                     *, api):
-        env = remove_python_envvars()
+                     *, api, env=None, ignore_stderr=False):
+        new_env = remove_python_envvars()
+        if env is not None:
+            new_env.update(env)
+        env = new_env
 
         if api == API_ISOLATED:
             default_preconfig = self.PRE_CONFIG_ISOLATED
@@ -634,7 +645,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
         out, err = self.run_embedded_interpreter(testname, env=env)
         if stderr is None and not expected_config['verbose']:
             stderr = ""
-        if stderr is not None:
+        if stderr is not None and not ignore_stderr:
             self.assertEqual(err.rstrip(), stderr)
         try:
             configs = json.loads(out)
@@ -966,6 +977,62 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
         self.check_all_configs("test_init_dont_parse_argv", config, pre_config,
                                api=API_PYTHON)
 
+    def test_init_setpath(self):
+        # Test Py_SetProgramName() + Py_SetPath()
+        config = self._get_expected_config()
+        paths = config['config']['module_search_paths']
+
+        config = {
+            'module_search_paths': paths,
+            'prefix': '',
+            'base_prefix': '',
+            'exec_prefix': '',
+            'base_exec_prefix': '',
+        }
+        env = {'TESTPATH': os.path.pathsep.join(paths)}
+        self.check_all_configs("test_init_setpath", config,
+                               api=API_COMPAT, env=env,
+                               ignore_stderr=True)
+
+    def test_init_setpythonhome(self):
+        # Test Py_SetPythonHome(home) + PYTHONPATH env var
+        # + Py_SetProgramName()
+        config = self._get_expected_config()
+        paths = config['config']['module_search_paths']
+        paths_str = os.path.pathsep.join(paths)
+
+        for path in paths:
+            if not os.path.isdir(path):
+                continue
+            if os.path.exists(os.path.join(path, 'os.py')):
+                home = os.path.dirname(path)
+                break
+        else:
+            self.fail(f"Unable to find home in {paths!r}")
+
+        prefix = exec_prefix = home
+        ver = sys.version_info
+        if MS_WINDOWS:
+            expected_paths = paths
+        else:
+            expected_paths = [
+                os.path.join(prefix, 'lib', f'python{ver.major}{ver.minor}.zip'),
+                os.path.join(home, 'lib', f'python{ver.major}.{ver.minor}'),
+                os.path.join(home, 'lib', f'python{ver.major}.{ver.minor}/lib-dynload')]
+
+        config = {
+            'home': home,
+            'module_search_paths': expected_paths,
+            'prefix': prefix,
+            'base_prefix': prefix,
+            'exec_prefix': exec_prefix,
+            'base_exec_prefix': exec_prefix,
+            'pythonpath_env': paths_str,
+        }
+        env = {'TESTHOME': home, 'TESTPATH': paths_str}
+        self.check_all_configs("test_init_setpythonhome", config,
+                               api=API_COMPAT, env=env)
+
 
 class AuditingTests(EmbeddingTestsMixin, unittest.TestCase):
     def test_open_code_hook(self):
diff --git a/Modules/getpath.c b/Modules/getpath.c
index 24e16b4..a4607ef 100644
--- a/Modules/getpath.c
+++ b/Modules/getpath.c
@@ -1012,12 +1012,12 @@ calculate_zip_path(PyCalculatePath *calculate, const wchar_t *prefix,
                    wchar_t *zip_path, size_t zip_path_len)
 {
     PyStatus status;
-    if (safe_wcscpy(zip_path, prefix, zip_path_len) < 0) {
-        return PATHLEN_ERR();
-    }
 
     if (calculate->prefix_found > 0) {
         /* Use the reduced prefix returned by Py_GetPrefix() */
+        if (safe_wcscpy(zip_path, prefix, zip_path_len) < 0) {
+            return PATHLEN_ERR();
+        }
         reduce(zip_path);
         reduce(zip_path);
     }
diff --git a/Programs/_testembed.c b/Programs/_testembed.c
index c3ccc0e..ed07606 100644
--- a/Programs/_testembed.c
+++ b/Programs/_testembed.c
@@ -20,11 +20,12 @@
  * Executed via 'EmbeddingTests' in Lib/test/test_capi.py
  *********************************************************/
 
+/* Use path starting with "./" avoids a search along the PATH */
+#define PROGRAM_NAME L"./_testembed"
+
 static void _testembed_Py_Initialize(void)
 {
-    /* HACK: the "./" at front avoids a search along the PATH in
-       Modules/getpath.c */
-    Py_SetProgramName(L"./_testembed");
+    Py_SetProgramName(PROGRAM_NAME);
     Py_Initialize();
 }
 
@@ -363,8 +364,7 @@ config_set_wide_string_list(PyConfig *config, PyWideStringList *list,
 
 static void config_set_program_name(PyConfig *config)
 {
-    /* Use path starting with "./" avoids a search along the PATH */
-    const wchar_t *program_name = L"./_testembed";
+    const wchar_t *program_name = PROGRAM_NAME;
     config_set_string(config, &config->program_name, program_name);
 }
 
@@ -1263,7 +1263,7 @@ static int _audit_hook_run(const char *eventName, PyObject *args, void *userData
 static int test_audit_run_command(void)
 {
     AuditRunCommandTest test = {"cpython.run_command"};
-    wchar_t *argv[] = {L"./_testembed", L"-c", L"pass"};
+    wchar_t *argv[] = {PROGRAM_NAME, L"-c", L"pass"};
 
     Py_IgnoreEnvironmentFlag = 0;
     PySys_AddAuditHook(_audit_hook_run, (void*)&test);
@@ -1274,7 +1274,7 @@ static int test_audit_run_command(void)
 static int test_audit_run_file(void)
 {
     AuditRunCommandTest test = {"cpython.run_file"};
-    wchar_t *argv[] = {L"./_testembed", L"filename.py"};
+    wchar_t *argv[] = {PROGRAM_NAME, L"filename.py"};
 
     Py_IgnoreEnvironmentFlag = 0;
     PySys_AddAuditHook(_audit_hook_run, (void*)&test);
@@ -1312,21 +1312,21 @@ static int run_audit_run_test(int argc, wchar_t **argv, void *test)
 static int test_audit_run_interactivehook(void)
 {
     AuditRunCommandTest test = {"cpython.run_interactivehook", 10};
-    wchar_t *argv[] = {L"./_testembed"};
+    wchar_t *argv[] = {PROGRAM_NAME};
     return run_audit_run_test(Py_ARRAY_LENGTH(argv), argv, &test);
 }
 
 static int test_audit_run_startup(void)
 {
     AuditRunCommandTest test = {"cpython.run_startup", 10};
-    wchar_t *argv[] = {L"./_testembed"};
+    wchar_t *argv[] = {PROGRAM_NAME};
     return run_audit_run_test(Py_ARRAY_LENGTH(argv), argv, &test);
 }
 
 static int test_audit_run_stdin(void)
 {
     AuditRunCommandTest test = {"cpython.run_stdin"};
-    wchar_t *argv[] = {L"./_testembed"};
+    wchar_t *argv[] = {PROGRAM_NAME};
     return run_audit_run_test(Py_ARRAY_LENGTH(argv), argv, &test);
 }
 
@@ -1423,6 +1423,88 @@ fail:
 }
 
 
+static int test_init_setpath(void)
+{
+    Py_SetProgramName(PROGRAM_NAME);
+
+    char *env = getenv("TESTPATH");
+    if (!env) {
+        fprintf(stderr, "missing TESTPATH env var\n");
+        return 1;
+    }
+    wchar_t *path = Py_DecodeLocale(env, NULL);
+    if (path == NULL) {
+        fprintf(stderr, "failed to decode TESTPATH\n");
+        return 1;
+    }
+    Py_SetPath(path);
+    PyMem_RawFree(path);
+    putenv("TESTPATH=");
+
+    Py_Initialize();
+    dump_config();
+    Py_Finalize();
+    return 0;
+}
+
+
+static int mysetenv(const char *name, const char *value)
+{
+    size_t len = strlen(name) + 1 + strlen(value) + 1;
+    char *env = PyMem_RawMalloc(len);
+    if (env == NULL) {
+        fprintf(stderr, "out of memory\n");
+        return -1;
+    }
+    strcpy(env, name);
+    strcat(env, "=");
+    strcat(env, value);
+
+    putenv(env);
+
+    /* Don't call PyMem_RawFree(env), but leak env memory block:
+       putenv() does not copy the string. */
+
+    return 0;
+}
+
+
+static int test_init_setpythonhome(void)
+{
+    char *env = getenv("TESTHOME");
+    if (!env) {
+        fprintf(stderr, "missing TESTHOME env var\n");
+        return 1;
+    }
+    wchar_t *home = Py_DecodeLocale(env, NULL);
+    if (home == NULL) {
+        fprintf(stderr, "failed to decode TESTHOME\n");
+        return 1;
+    }
+    Py_SetPythonHome(home);
+    PyMem_RawFree(home);
+    putenv("TESTHOME=");
+
+    char *path = getenv("TESTPATH");
+    if (!path) {
+        fprintf(stderr, "missing TESTPATH env var\n");
+        return 1;
+    }
+
+    if (mysetenv("PYTHONPATH", path) < 0) {
+        return 1;
+    }
+    putenv("TESTPATH=");
+
+    Py_SetProgramName(PROGRAM_NAME);
+
+    Py_Initialize();
+    dump_config();
+    Py_Finalize();
+    return 0;
+}
+
+
 static void configure_init_main(PyConfig *config)
 {
     wchar_t* argv[] = {
@@ -1559,7 +1641,10 @@ static struct TestCase TestCases[] = {
     {"test_init_run_main", test_init_run_main},
     {"test_init_main", test_init_main},
     {"test_init_sys_add", test_init_sys_add},
+    {"test_init_setpath", test_init_setpath},
+    {"test_init_setpythonhome", test_init_setpythonhome},
     {"test_run_main", test_run_main},
+
     {"test_open_code_hook", test_open_code_hook},
     {"test_audit", test_audit},
     {"test_audit_subinterpreter", test_audit_subinterpreter},