/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.flink.client.python; import org.apache.flink.configuration.Configuration; import org.apache.flink.core.fs.Path; import org.apache.flink.core.testutils.CommonTestUtils; import org.apache.flink.util.FileUtils; import org.apache.flink.util.OperatingSystem; import org.junit.After; import org.junit.Assert; import org.junit.Assume; import org.junit.Before; import org.junit.Test; import java.io.File; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Paths; import java.nio.file.StandardOpenOption; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; import java.util.UUID; import java.util.stream.Collectors; import static org.apache.flink.client.python.PythonEnvUtils.PYFLINK_CLIENT_EXECUTABLE; import static org.apache.flink.client.python.PythonEnvUtils.preparePythonEnvironment; import static org.apache.flink.python.PythonOptions.PYTHON_CLIENT_EXECUTABLE; import static org.apache.flink.python.PythonOptions.PYTHON_FILES; import static org.apache.flink.python.util.PythonDependencyUtils.FILE_DELIMITER; /** * Tests for the {@link PythonEnvUtils}. */ public class PythonEnvUtilsTest { private String tmpDirPath; @Before public void prepareTestEnvironment() { File tmpDirFile = new File(System.getProperty("java.io.tmpdir"), "pyflink_" + UUID.randomUUID()); tmpDirFile.mkdirs(); this.tmpDirPath = tmpDirFile.getAbsolutePath(); } @Test public void testPreparePythonEnvironment() throws IOException { // Skip this test on Windows as we can not control the Window Driver letters. Assume.assumeFalse(OperatingSystem.isWindows()); // xxx/a.zip, xxx/subdir/b.py, xxx/subdir/c.zip File zipFile = new File(tmpDirPath + File.separator + "a.zip"); File dirFile = new File(tmpDirPath + File.separator + "module_dir"); File subdirFile = new File(tmpDirPath + File.separator + "subdir"); File relativeFile = new File(tmpDirPath + File.separator + "subdir" + File.separator + "b.py"); File schemeFile = new File(tmpDirPath + File.separator + "subdir" + File.separator + "c.zip"); // The files must actually exist zipFile.createNewFile(); dirFile.mkdir(); subdirFile.mkdir(); relativeFile.createNewFile(); schemeFile.createNewFile(); String workingDir = new File("").getAbsolutePath(); String absolutePath = relativeFile.getAbsolutePath(); Path zipPath = new Path(zipFile.getAbsolutePath()); Path dirPath = new Path(dirFile.getAbsolutePath()); Path relativePath = new Path(Paths.get(workingDir).relativize(Paths.get(absolutePath)).toString()); Path schemePath = new Path("file://" + schemeFile.getAbsolutePath()); List<Path> pyFilesList = new ArrayList<>(); pyFilesList.add(zipPath); pyFilesList.add(dirPath); pyFilesList.add(relativePath); pyFilesList.add(schemePath); String pyFiles = pyFilesList.stream() .map(Path::toString) .collect(Collectors.joining(FILE_DELIMITER)); Configuration config = new Configuration(); config.set(PYTHON_FILES, pyFiles); PythonEnvUtils.PythonEnvironment env = preparePythonEnvironment(config, null, tmpDirPath); String base = replaceUUID(env.tempDirectory); Set<String> expectedPythonPaths = new HashSet<>(); expectedPythonPaths.add(String.join(File.separator, base, "{uuid}", "a.zip")); expectedPythonPaths.add(String.join(File.separator, base, "{uuid}", "module_dir")); expectedPythonPaths.add(String.join(File.separator, base, "{uuid}")); expectedPythonPaths.add(String.join(File.separator, base, "{uuid}", "c.zip")); Set<String> actualPaths = Arrays.stream(env.pythonPath.split(File.pathSeparator)) .map(PythonEnvUtilsTest::replaceUUID) .collect(Collectors.toSet()); Assert.assertEquals(expectedPythonPaths, actualPaths); } @Test public void testStartPythonProcess() { PythonEnvUtils.PythonEnvironment pythonEnv = new PythonEnvUtils.PythonEnvironment(); pythonEnv.tempDirectory = tmpDirPath; pythonEnv.pythonPath = tmpDirPath; List<String> commands = new ArrayList<>(); String pyPath = String.join(File.separator, tmpDirPath, "verifier.py"); try { File pyFile = new File(pyPath); pyFile.createNewFile(); pyFile.setExecutable(true); String pyProgram = "#!/usr/bin/python\n" + "# -*- coding: UTF-8 -*-\n" + "import os\n" + "import sys\n" + "\n" + "if __name__=='__main__':\n" + "\tfilename = sys.argv[1]\n" + "\tfo = open(filename, \"w\")\n" + "\tfo.write(os.getcwd())\n" + "\tfo.close()"; Files.write(pyFile.toPath(), pyProgram.getBytes(), StandardOpenOption.WRITE); String result = String.join(File.separator, tmpDirPath, "python_working_directory.txt"); commands.add(pyPath); commands.add(result); Process pythonProcess = PythonEnvUtils.startPythonProcess(pythonEnv, commands); int exitCode = pythonProcess.waitFor(); if (exitCode != 0) { throw new RuntimeException("Python process exits with code: " + exitCode); } String cmdResult = new String(Files.readAllBytes(new File(result).toPath())); // Check if the working directory of python process is the same as java process. Assert.assertEquals(cmdResult, System.getProperty("user.dir")); pythonProcess.destroyForcibly(); pyFile.delete(); new File(result).delete(); } catch (IOException | InterruptedException e) { throw new RuntimeException("test start Python process failed " + e.getMessage()); } } @Test public void testSetPythonExecutable() throws IOException { Configuration config = new Configuration(); PythonEnvUtils.PythonEnvironment env = preparePythonEnvironment(config, null, tmpDirPath); if (OperatingSystem.isWindows()) { Assert.assertEquals("python.exe", env.pythonExec); } else { Assert.assertEquals("python", env.pythonExec); } Map<String, String> systemEnv = new HashMap<>(System.getenv()); systemEnv.put(PYFLINK_CLIENT_EXECUTABLE, "python3"); CommonTestUtils.setEnv(systemEnv); try { env = preparePythonEnvironment(config, null, tmpDirPath); Assert.assertEquals("python3", env.pythonExec); } finally { systemEnv.remove(PYFLINK_CLIENT_EXECUTABLE); CommonTestUtils.setEnv(systemEnv); } config.set(PYTHON_CLIENT_EXECUTABLE, "/usr/bin/python"); env = preparePythonEnvironment(config, null, tmpDirPath); Assert.assertEquals("/usr/bin/python", env.pythonExec); } @Test public void testPrepareEnvironmentWithEntryPointScript() throws IOException { File entryFile = new File(tmpDirPath + File.separator + "test.py"); // The file must actually exist entryFile.createNewFile(); String entryFilePath = entryFile.getAbsolutePath(); Configuration config = new Configuration(); PythonEnvUtils.PythonEnvironment env = preparePythonEnvironment(config, entryFilePath, tmpDirPath); Set<String> expectedPythonPaths = new HashSet<>(); expectedPythonPaths.add( new Path(String.join(File.separator, replaceUUID(env.tempDirectory), "{uuid}")).toString()); Set<String> actualPaths = Arrays.stream(env.pythonPath.split(File.pathSeparator)) .map(PythonEnvUtilsTest::replaceUUID) .collect(Collectors.toSet()); Assert.assertEquals(expectedPythonPaths, actualPaths); } @After public void cleanEnvironment() { FileUtils.deleteDirectoryQuietly(new File(tmpDirPath)); } private static String replaceUUID(String originPath) { return originPath.replaceAll("[a-z0-9]{8}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{12}", "{uuid}"); } }