import io import os import sys import unittest from unittest.mock import MagicMock, patch import fastdeploy.collect_env as collect_env class TestCollectEnv(unittest.TestCase): def setUp(self): self.run_lambda = MagicMock() self.run_lambda.return_value = (0, "test output", "") def test_run(self): result = collect_env.run("echo test") self.assertIsInstance(result, tuple) def test_run_nvidia(self): result = collect_env.run("nvidia-smi topo -m") self.assertIsInstance(result, tuple) def test_run_and_read_all(self): result = collect_env.run_and_read_all(self.run_lambda, "test command") self.assertEqual(result, "test output") self.run_lambda.return_value = (1, "version 1.0", "") result = collect_env.run_and_read_all(self.run_lambda, "test command") self.assertEqual(result, None) def test_run_and_parse_first_match(self): self.run_lambda.return_value = (0, "version 1.0", "") result = collect_env.run_and_parse_first_match(self.run_lambda, "test command", r"version (.*)") self.assertEqual(result, "1.0") self.run_lambda.return_value = (1, "version 1.0", "") result = collect_env.run_and_parse_first_match(self.run_lambda, "test command", r"version (.*)") self.assertEqual(result, None) self.run_lambda.return_value = (0, "version 1.0", "") result = collect_env.run_and_parse_first_match(self.run_lambda, "test command", r"sadsad") self.assertEqual(result, None) def test_run_and_return_first_line(self): self.run_lambda.return_value = (0, "line1\nline2", "") result = collect_env.run_and_return_first_line(self.run_lambda, "test command") self.assertEqual(result, "line1") self.run_lambda.return_value = (1, "line1\nline2", "") result = collect_env.run_and_return_first_line(self.run_lambda, "test command") self.assertEqual(result, None) def test_get_conda_packages(self): with patch("fastdeploy.collect_env.run_and_read_all") as mock_read: mock_read.return_value = "package1\npackage2" result = collect_env.get_conda_packages(self.run_lambda) self.assertIsNotNone(result) with patch("fastdeploy.collect_env.run_and_read_all") as mock_read: mock_read.return_value = None result = collect_env.get_conda_packages(self.run_lambda) self.assertIsNone(result) def test_get_gcc_version(self): with patch("fastdeploy.collect_env.run_and_parse_first_match") as mock_parse: mock_parse.return_value = "1.0" result = collect_env.get_gcc_version(self.run_lambda) self.assertEqual(result, "1.0") def test_get_clang_version(self): with patch("fastdeploy.collect_env.run_and_parse_first_match") as mock_parse: mock_parse.return_value = "1.0" result = collect_env.get_clang_version(self.run_lambda) self.assertEqual(result, "1.0") def test_get_cmake_version(self): with patch("fastdeploy.collect_env.run_and_parse_first_match") as mock_parse: mock_parse.return_value = "1.0" result = collect_env.get_cmake_version(self.run_lambda) self.assertEqual(result, "1.0") def test_get_nvidia_driver_version(self): with patch("fastdeploy.collect_env.run_and_parse_first_match") as mock_parse: mock_parse.return_value = "1.0" result = collect_env.get_nvidia_driver_version(self.run_lambda) self.assertEqual(result, "1.0") with ( patch("fastdeploy.collect_env.run_and_parse_first_match", return_value="1.0"), patch("fastdeploy.collect_env.get_platform", return_value="darwin"), ): result = collect_env.get_nvidia_driver_version(self.run_lambda) self.assertEqual(result, "1.0") def test_get_gpu_info(self): with patch("fastdeploy.collect_env.TORCH_AVAILABLE", False): result = collect_env.get_gpu_info(self.run_lambda) self.assertIsNotNone(result) with ( patch("fastdeploy.collect_env.get_platform", return_value="darwin"), patch("fastdeploy.collect_env.TORCH_AVAILABLE", True), patch("fastdeploy.collect_env.torch", create=True), ): result = collect_env.get_gpu_info(self.run_lambda) self.assertIsNotNone(result) def test_get_running_cuda_version(self): with patch("fastdeploy.collect_env.run_and_parse_first_match") as mock_parse: mock_parse.return_value = "1.0" result = collect_env.get_running_cuda_version(self.run_lambda) self.assertEqual(result, "1.0") def test_get_cudnn_version(self): with ( patch("fastdeploy.collect_env.run") as mock_run, patch("fastdeploy.collect_env.get_platform", return_value="linux"), ): mock_run.return_value = (0, "/usr/local/cuda/lib64/libcudnn.so.8.4.1", "") result = collect_env.get_cudnn_version(self.run_lambda) self.assertEqual(result, None) with ( patch("fastdeploy.collect_env.run") as mock_run, patch("fastdeploy.collect_env.get_platform", return_value="win32"), ): mock_run.return_value = (0, "/usr/local/cuda/lib64/libcudnn.so.8.4.1", "") result = collect_env.get_cudnn_version(self.run_lambda) self.assertEqual(result, None) with ( patch("fastdeploy.collect_env.run") as mock_run, patch("fastdeploy.collect_env.get_platform", return_value="darwin"), ): mock_run.return_value = (0, "/usr/local/cuda/lib64/libcudnn.so.8.4.1", "") result = collect_env.get_cudnn_version(self.run_lambda) self.assertEqual(result, None) with ( patch("fastdeploy.collect_env.run") as mock_run, patch("fastdeploy.collect_env.get_platform", return_value="darwin"), ): mock_run.return_value = (2, "/usr/local/cuda/lib64/libcudnn.so.8.4.1\n/usr/xxx", "") self.run_lambda.return_value = (2, "version 1.0", "") result = collect_env.get_cudnn_version(self.run_lambda) self.assertEqual(result, None) with ( patch("os.path.realpath", side_effect=lambda x: f"/real_path/to/{x.split('/')[-1]}"), patch("os.path.isfile", return_value=True), ): self.run_lambda.return_value = ( 0, "/usr/local/cuda/lib/libcudnn.so.8\n/usr/local/cuda/lib/libcudnn.so.8.2.1", "", ) cudnn_version = collect_env.get_cudnn_version(self.run_lambda) # 验证返回结果是预期的多行字符串 expected_output = ( "Probably one of the following:\n/real_path/to/libcudnn.so.8\n/real_path/to/libcudnn.so.8.2.1" ) self.assertEqual(cudnn_version, expected_output) def test_get_nvidia_smi(self): result = collect_env.get_nvidia_smi() self.assertIsNotNone(result) with patch("fastdeploy.collect_env.get_platform", return_value="win32"): result = collect_env.get_nvidia_smi() self.assertIsNotNone(result) def test_get_fastdeploy_version(self): with patch("fastdeploy.collect_env.os.environ.get", return_value="1.0"): result = collect_env.get_fastdeploy_version() self.assertEqual(result, "1.0") with patch("fastdeploy.collect_env.os.environ.get", return_value=None): result = collect_env.get_fastdeploy_version() self.assertIsNotNone(result) with patch("pkg_resources.get_distribution", side_effect=Exception("Package not found")): with patch("fastdeploy.collect_env.os.environ.get", return_value=None): with patch("subprocess.run", return_value=None): result = collect_env.get_fastdeploy_version() self.assertIsNotNone(result) def test_summarize_fastdeploy_build_flags(self): result = collect_env.summarize_fastdeploy_build_flags() self.assertIsNotNone(result) def test_get_gpu_topo(self): result = collect_env.get_gpu_topo(self.run_lambda) self.assertIsNotNone(result) def test_get_cpu_info(self): self.run_lambda.return_value = (0, "Architecture: x86_64\nModel name: Intel(R) Xeon(R) CPU", "") with patch("fastdeploy.collect_env.get_platform", return_value="linux"): cpu_info = collect_env.get_cpu_info(self.run_lambda) self.assertIn("x86_64", cpu_info) self.assertIn("Intel(R) Xeon(R)", cpu_info) with patch("fastdeploy.collect_env.get_platform", return_value="win32"): cpu_info = collect_env.get_cpu_info(self.run_lambda) self.assertIn("x86_64", cpu_info) self.assertIn("Intel(R) Xeon(R)", cpu_info) with patch("fastdeploy.collect_env.get_platform", return_value="darwin"): cpu_info = collect_env.get_cpu_info(self.run_lambda) self.assertIn("x86_64", cpu_info) self.assertIn("Intel(R) Xeon(R)", cpu_info) self.run_lambda.return_value = (1, "Architecture: x86_64\nModel name: Intel(R) Xeon(R) CPU", "err") with patch("fastdeploy.collect_env.get_platform", return_value="darwin"): cpu_info = collect_env.get_cpu_info(self.run_lambda) self.assertIn("err", cpu_info) def test_get_platform(self): with patch("sys.platform", "linux"): self.assertEqual(collect_env.get_platform(), "linux") with patch("sys.platform", "win32"): self.assertEqual(collect_env.get_platform(), "win32") with patch("sys.platform", "cygwin"): self.assertEqual(collect_env.get_platform(), "cygwin") with patch("sys.platform", "darwin"): self.assertEqual(collect_env.get_platform(), "darwin") def test_get_os_linux_lsb_success(self): """测试 Linux 环境下，lsb_release 命令成功。""" with patch("sys.platform", "linux"): # 模拟 get_lsb_version 成功返回 with patch("fastdeploy.collect_env.get_lsb_version", return_value="Ubuntu 20.04 LTS"): # 模拟 platform.machine 成功返回 with patch("platform.machine", return_value="x86_64"): result = collect_env.get_os(self.run_lambda) self.assertEqual(result, "Ubuntu 20.04 LTS (x86_64)") def test_get_os_linux_lsb_fail_check_release_success(self): """测试 Linux 环境下，lsb_release 失败，但 /etc/*-release 成功。""" with patch("sys.platform", "linux"): # 模拟 get_lsb_version 失败 with patch("fastdeploy.collect_env.get_lsb_version", return_value=None): # 模拟 check_release_file 成功返回 with patch("fastdeploy.collect_env.check_release_file", return_value="CentOS Linux 8"): with patch("platform.machine", return_value="x86_64"): result = collect_env.get_os(self.run_lambda) self.assertEqual(result, "CentOS Linux 8 (x86_64)") def test_get_os_linux_all_fail(self): """测试 Linux 环境下，所有方法都失败。""" with patch("sys.platform", "linux"): with patch("fastdeploy.collect_env.get_lsb_version", return_value=None): with patch("fastdeploy.collect_env.check_release_file", return_value=None): with patch("platform.machine", return_value="x86_64"): result = collect_env.get_os(self.run_lambda) self.assertEqual(result, "linux (x86_64)") def test_get_os_windows_success(self): """测试 Windows 环境下，命令成功。""" with patch("sys.platform", "win32"): # 模拟 get_windows_version 成功返回 with patch("fastdeploy.collect_env.get_windows_version", return_value="Microsoft Windows 10"): result = collect_env.get_os(self.run_lambda) self.assertEqual(result, "Microsoft Windows 10") def test_get_os_windows_fail(self): """测试 Windows 环境下，命令失败。""" with patch("sys.platform", "win32"): # 模拟 get_windows_version 失败返回 None with patch("fastdeploy.collect_env.get_windows_version", return_value=None): result = collect_env.get_os(self.run_lambda) self.assertIsNone(result) def test_get_os_macos_success(self): """测试 macOS 环境下，命令成功。""" with patch("sys.platform", "darwin"): # 模拟 get_mac_version 成功返回 with patch("fastdeploy.collect_env.get_mac_version", return_value="12.3.1"): with patch("platform.machine", return_value="arm64"): result = collect_env.get_os(self.run_lambda) self.assertEqual(result, "macOS 12.3.1 (arm64)") def test_get_os_macos_fail(self): """测试 macOS 环境下，命令失败。""" with patch("sys.platform", "darwin"): # 模拟 get_mac_version 失败返回 None with patch("fastdeploy.collect_env.get_mac_version", return_value=None): result = collect_env.get_os(self.run_lambda) self.assertIsNone(result) def test_get_os_unknown_platform(self): """测试未知平台的情况。""" with patch("sys.platform", "solaris"): result = collect_env.get_os(self.run_lambda) self.assertEqual(result, "solaris") def test_get_python_platform(self): """测试 get_python_platform 函数返回正确值。""" with patch("platform.platform", return_value="Linux-5.15.0-76-generic-x86_64-with-glibc2.35"): result = collect_env.get_python_platform() self.assertEqual(result, "Linux-5.15.0-76-generic-x86_64-with-glibc2.35") def test_get_libc_version_linux_success(self): """测试在 Linux 环境下成功获取 libc 版本。""" with patch("fastdeploy.collect_env.get_platform", return_value="linux"): with patch("platform.libc_ver", return_value=("glibc", "2.35")): result = collect_env.get_libc_version() self.assertEqual(result, "glibc-2.35") def test_get_libc_version_non_linux(self): """测试在非 Linux 环境下返回 'N/A'。""" with patch("fastdeploy.collect_env.get_platform", return_value="win32"): # 确保 platform.libc_ver() 不被调用，或者即使调用也不会影响结果 with patch("platform.libc_ver") as mock_libc_ver: result = collect_env.get_libc_version() self.assertEqual(result, "N/A") mock_libc_ver.assert_not_called() def test_get_pip_packages_no_pip_or_uv(self): """Test that a RuntimeError is raised when neither pip nor uv are available.""" with patch.dict("os.environ", {}, clear=True), patch("importlib.util.find_spec", return_value=None): with self.assertRaises(RuntimeError) as cm: collect_env.get_pip_packages(self.run_lambda) self.assertIn("Could not collect pip list output", str(cm.exception)) def test_get_pip_packages_success(self): """Test that the pip module is available and a list of packages is returned.""" with ( patch("fastdeploy.collect_env.run_and_read_all") as mock_run_and_read_all, patch("importlib.util.find_spec", return_value=True), ): mock_run_and_read_all.return_value = "torch==2.0.0\nregex==2023.1.1\nnumpy==1.25.0" pip_version, packages = collect_env.get_pip_packages(self.run_lambda) self.assertEqual(pip_version, "pip3" if sys.version[0] == "3" else "pip") self.assertIn("torch==2.0.0", packages) self.assertIn("numpy==1.25.0", packages) self.assertNotIn("regex", packages) def test_get_pip_packages_uv_available(self): """Test that uv is used when pip is not available but the UV environment variable is set.""" with ( patch.dict("os.environ", {"UV": "1"}), patch("fastdeploy.collect_env.run_and_read_all") as mock_run_and_read_all, patch("importlib.util.find_spec", return_value=False), ): mock_run_and_read_all.return_value = "torch==2.0.0\nregex==2023.1.1\nnumpy==1.25.0" pip_version, packages = collect_env.get_pip_packages(self.run_lambda) self.assertIsNotNone(packages) def test_get_pip_packages_command_fail(self): """Test that an empty string is returned when the pip command fails.""" with ( patch("fastdeploy.collect_env.run_and_read_all", return_value="\n"), patch("importlib.util.find_spec", return_value=True), ): pip_version, packages = collect_env.get_pip_packages(self.run_lambda) self.assertEqual(packages, "") def test_get_pip_packages_custom_patterns(self): """Test that the function correctly filters packages based on custom patterns.""" with ( patch("fastdeploy.collect_env.run_and_read_all") as mock_run_and_read_all, patch("importlib.util.find_spec", return_value=True), ): mock_run_and_read_all.return_value = "torch==2.0.0\nmy-custom-lib==1.0.0\nrequests==2.28.1" custom_patterns = {"my-custom-lib", "requests"} pip_version, packages = collect_env.get_pip_packages(self.run_lambda, patterns=custom_patterns) self.assertIn("my-custom-lib==1.0.0", packages) self.assertIn("requests==2.28.1", packages) self.assertNotIn("torch", packages) def test_xnnpack_available_with_torch(self): with patch("fastdeploy.collect_env.TORCH_AVAILABLE", True), patch("fastdeploy.collect_env.torch", create=True): result = collect_env.is_xnnpack_available() self.assertIsNotNone(result) def test_xnnpack_not_available_without_torch(self): """测试 torch 不可用时，返回 'N/A'。""" with patch("fastdeploy.collect_env.TORCH_AVAILABLE", False): result = collect_env.is_xnnpack_available() self.assertEqual(result, "N/A") def test_get_env_vars_with_relevant_vars(self): """测试正确收集相关环境变量。""" # 准备一个包含各种类型环境变量的字典 mock_env = { "TORCH_DEBUG": "1", "CUDA_ARCHS": "7.5", "SOME_OTHER_VAR": "value", "MY_API_KEY": "secret_key", "FASTDEPLOY_MODEL_DIR": "/path/to/model", # 假设这个在 environment_variables 中 } # 模拟 environment_variables 列表 with patch.dict(os.environ, mock_env, clear=True): with patch( "fastdeploy.collect_env.environment_variables", ["FASTDEPLOY_MODEL_DIR"], ): env_vars_string = collect_env.get_env_vars() self.assertIn("TORCH_DEBUG=1", env_vars_string) self.assertIn("CUDA_ARCHS=7.5", env_vars_string) self.assertIn("FASTDEPLOY_MODEL_DIR=/path/to/model", env_vars_string) self.assertNotIn("SOME_OTHER_VAR", env_vars_string) self.assertNotIn("MY_API_KEY", env_vars_string) def test_get_cuda_config_with_both_vars_set(self): with patch("fastdeploy.collect_env.TORCH_AVAILABLE", True), patch("fastdeploy.collect_env.torch", create=True): mock_env = { "CUDA_MODULE_LOADING": "xxx", } with patch.dict(os.environ, mock_env, clear=True): result = collect_env.get_cuda_module_loading_config() self.assertEqual(result, "xxx") def test_get_cuda_config_with_no_vars_set(self): """测试两个环境变量都未设置。""" with patch("fastdeploy.collect_env.TORCH_AVAILABLE", False): result = collect_env.get_cuda_module_loading_config() self.assertEqual(result, "N/A") def test_get_env_info_full(self): # 使用嵌套的 with patch 语句来模拟所有依赖函数的返回值 with ( patch("fastdeploy.collect_env.get_platform", return_value="linux"), patch("fastdeploy.collect_env.get_os", return_value="Ubuntu 20.04 (x86_64)"), patch("fastdeploy.collect_env.get_python_platform", return_value="Python 3.8.10"), patch("fastdeploy.collect_env.get_cuda_module_loading_config", return_value="CUDA_DEVICE_VISIBLE=0"), patch("fastdeploy.collect_env.get_libc_version", return_value="glibc-2.31"), patch("fastdeploy.collect_env.get_gcc_version", return_value="9.3.0"), patch("fastdeploy.collect_env.get_clang_version", return_value=None), patch("fastdeploy.collect_env.get_cmake_version", return_value="3.18.4"), patch("fastdeploy.collect_env.get_nvidia_driver_version", return_value="470.82.00"), patch("fastdeploy.collect_env.get_cudnn_version", return_value="8.2.1"), patch("fastdeploy.collect_env.get_running_cuda_version", return_value="11.4"), patch("fastdeploy.collect_env.get_gpu_info", return_value="GeForce RTX 3080"), patch("fastdeploy.collect_env.get_cpu_info", return_value="Intel(R) Core(TM) i9-10900K"), patch("fastdeploy.collect_env.is_xnnpack_available", return_value="True"), patch("fastdeploy.collect_env.get_fastdeploy_version", return_value="1.0.0"), patch("fastdeploy.collect_env.get_conda_packages", return_value="numpy==1.22.0\ntorch==1.11.0"), patch("fastdeploy.collect_env.get_pip_packages", return_value=("pip3", "requests==2.27.1\nscipy==1.8.0")), patch("fastdeploy.collect_env.get_env_vars", return_value="CUDA_VISIBLE_DEVICES=0"), patch("fastdeploy.collect_env.get_gpu_topo", return_value="GPU-Direct disabled"), patch("fastdeploy.collect_env.torch", create=True, __version__="1.0.0"), ): info_string = collect_env.get_env_info() self.assertIsNotNone(info_string) def test_get_env_info_all_na(self): with ( patch("fastdeploy.collect_env.get_python_platform", return_value="N/A"), patch("fastdeploy.collect_env.get_cuda_module_loading_config", return_value="N/A"), patch("fastdeploy.collect_env.get_libc_version", return_value="N/A"), patch("fastdeploy.collect_env.get_gcc_version", return_value="N/A"), patch("fastdeploy.collect_env.get_clang_version", return_value="N/A"), patch("fastdeploy.collect_env.get_cmake_version", return_value="N/A"), patch("fastdeploy.collect_env.get_nvidia_driver_version", return_value="N/A"), patch("fastdeploy.collect_env.get_cudnn_version", return_value="N/A"), patch("fastdeploy.collect_env.get_running_cuda_version", return_value="N/A"), patch("fastdeploy.collect_env.get_gpu_info", return_value="N/A"), patch("fastdeploy.collect_env.get_cpu_info", return_value="N/A"), patch("fastdeploy.collect_env.is_xnnpack_available", return_value="N/A"), patch("fastdeploy.collect_env.get_fastdeploy_version", return_value="N/A"), patch("fastdeploy.collect_env.get_conda_packages", return_value=None), patch("fastdeploy.collect_env.get_pip_packages", return_value=("N/A", "N/A")), patch("fastdeploy.collect_env.get_env_vars", return_value=""), patch("fastdeploy.collect_env.get_gpu_topo", return_value="N/A"), patch("fastdeploy.collect_env.TORCH_AVAILABLE", return_value=False), patch("fastdeploy.collect_env.PADDLE_AVAILABLE", return_value=False), patch("fastdeploy.collect_env.torch", create=True, __version__="1.0.0"), ): info_string = collect_env.get_env_info() self.assertIsNotNone(info_string) def test_main_with_collect(self): captured_output = io.StringIO() with ( patch("sys.stdout", new=captured_output), patch("fastdeploy.collect_env.torch", create=True, __version__="1.0.0"), patch("fastdeploy.collect_env.hasattr", return_value=False), ): collect_env.main() output = captured_output.getvalue() expected_message = "Collecting environment information" self.assertIn(expected_message, output) if __name__ == "__main__": unittest.main()