Line data    Source code

       1              : #include "backend/engine/semantic/python_udf_runtime.h"
       2              : 
       3              : #include <cstdlib>
       4              : #include <filesystem>
       5              : #include <fstream>
       6              : #include <string>
       7              : 
       8              : #include "gtest/gtest.h"
       9              : 
      10              : namespace bigquery_emulator {
      11              : namespace backend {
      12              : namespace engine {
      13              : namespace semantic {
      14              : namespace {
      15              : 
      16              : namespace fs = std::filesystem;
      17              : 
      18              : class ScopedEnv {
      19              :  public:
      20            3 :   ScopedEnv(const char* key, const char* value) : key_(key) {
      21            3 :     if (const char* prev = std::getenv(key)) {
      22            0 :       had_prev_ = true;
      23            0 :       prev_ = prev;
      24            0 :     }
      25            3 :     setenv(key, value, 1);
      26            3 :   }
      27            3 :   ~ScopedEnv() {
      28            3 :     if (had_prev_) {
      29            0 :       setenv(key_, prev_.c_str(), 1);
      30            3 :     } else {
      31            3 :       unsetenv(key_);
      32            3 :     }
      33            3 :   }
      34              : 
      35              :  private:
      36              :   const char* key_;
      37              :   bool had_prev_ = false;
      38              :   std::string prev_;
      39              : };
      40              : 
      41              : }  // namespace
      42              : 
      43            1 : TEST(PythonUdfRuntimeTest, ImportModuleNameFromPackageSpecStripsVersions) {
      44            1 :   EXPECT_EQ(ImportModuleNameFromPackageSpec("lxml"), "lxml");
      45            1 :   EXPECT_EQ(ImportModuleNameFromPackageSpec("lxml==4.9.3"), "lxml");
      46            1 :   EXPECT_EQ(ImportModuleNameFromPackageSpec("pandas>=2.1"), "pandas");
      47            1 :   EXPECT_EQ(ImportModuleNameFromPackageSpec(" google-cloud-translate==3.11 "),
      48            1 :             "google-cloud-translate");
      49            1 : }
      50              : 
      51            1 : TEST(PythonUdfRuntimeTest, ResolvePythonInterpreterPathPrefersExplicitEnv) {
      52            1 :   ScopedEnv python("BIGQUERY_EMULATOR_PYTHON", "/tmp/bqemu-test-python");
      53            1 :   ScopedEnv data_dir("BIGQUERY_EMULATOR_DATA_DIR", "/tmp/bqemu-data");
      54            1 :   auto path_or = ResolvePythonInterpreterPath();
      55            1 :   ASSERT_TRUE(path_or.ok());
      56            1 :   EXPECT_EQ(*path_or, "/tmp/bqemu-test-python");
      57            1 : }
      58              : 
      59            1 : TEST(PythonUdfRuntimeTest, ResolvePythonInterpreterPathUsesManagedVenv) {
      60            1 :   const fs::path root = fs::temp_directory_path() / "bqemu-managed-venv-test";
      61            1 :   fs::remove_all(root);
      62            1 :   const fs::path python = root / "python-udf-env" / "bin" / "python3";
      63            1 :   ASSERT_TRUE(fs::create_directories(python.parent_path()));
      64            1 :   {
      65            1 :     std::ofstream out(python);
      66            1 :     out << "#!/bin/sh\nexit 0\n";
      67            1 :   }
      68            1 :   fs::permissions(python,
      69            1 :                   fs::perms::owner_all | fs::perms::group_read |
      70            1 :                       fs::perms::group_exec | fs::perms::others_read |
      71            1 :                       fs::perms::others_exec);
      72            1 :   ScopedEnv data_dir("BIGQUERY_EMULATOR_DATA_DIR", root.string().c_str());
      73            1 :   unsetenv("BIGQUERY_EMULATOR_PYTHON");
      74            1 :   auto path_or = ResolvePythonInterpreterPath();
      75            1 :   ASSERT_TRUE(path_or.ok());
      76            1 :   EXPECT_EQ(*path_or, python.string());
      77            1 :   fs::remove_all(root);
      78            1 : }
      79              : 
      80            1 : TEST(PythonUdfRuntimeTest, ResolvePythonInterpreterPathFallsBackToHost) {
      81            1 :   unsetenv("BIGQUERY_EMULATOR_PYTHON");
      82            1 :   unsetenv("BIGQUERY_EMULATOR_DATA_DIR");
      83            1 :   auto path_or = ResolvePythonInterpreterPath();
      84            1 :   ASSERT_TRUE(path_or.ok());
      85            1 :   EXPECT_EQ(*path_or, "python3");
      86            1 : }
      87              : 
      88              : }  // namespace semantic
      89              : }  // namespace engine
      90              : }  // namespace backend
      91              : }  // namespace bigquery_emulator