refactor: simplify the plugin detection code, add test coverage (#1769)

shobsi · web-flow · commit 9df26ea6e585 · 2025-05-29T12:47:05.000-07:00
We added instrumentation of the GCP vscode extension and jupyter plugin in pandas-gbq and bigquery-magics. In this change we are propagating the improvements there to BigFrames.
diff --git a/packages/bigframes/bigframes/session/environment.py b/packages/bigframes/bigframes/session/environment.py
@@ -1,4 +1,4 @@
-# Copyright 2023 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,9 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 import importlib
 import json
 import os
+import pathlib
+
+Path = pathlib.Path
+
 
 # The identifier for GCP VS Code extension
 # https://cloud.google.com/code/docs/vscode/install
@@ -29,40 +34,36 @@
 def _is_vscode_extension_installed(extension_id: str) -> bool:
     """
     Checks if a given Visual Studio Code extension is installed.
-
     Args:
         extension_id: The ID of the extension (e.g., "ms-python.python").
-
     Returns:
         True if the extension is installed, False otherwise.
     """
     try:
         # Determine the user's VS Code extensions directory.
-        user_home = os.path.expanduser("~")
-        if os.name == "nt":  # Windows
-            vscode_extensions_dir = os.path.join(user_home, ".vscode", "extensions")
-        elif os.name == "posix":  # macOS and Linux
-            vscode_extensions_dir = os.path.join(user_home, ".vscode", "extensions")
-        else:
-            raise OSError("Unsupported operating system.")
+        user_home = Path.home()
+        vscode_extensions_dir = user_home / ".vscode" / "extensions"
 
         # Check if the extensions directory exists.
-        if os.path.exists(vscode_extensions_dir):
-            # Iterate through the subdirectories in the extensions directory.
-            for item in os.listdir(vscode_extensions_dir):
-                item_path = os.path.join(vscode_extensions_dir, item)
-                if os.path.isdir(item_path) and item.startswith(extension_id + "-"):
-                    # Check if the folder starts with the extension ID.
-                    # Further check for manifest file, as a more robust check.
-                    manifest_path = os.path.join(item_path, "package.json")
-                    if os.path.exists(manifest_path):
-                        try:
-                            with open(manifest_path, "r", encoding="utf-8") as f:
-                                json.load(f)
-                            return True
-                        except (FileNotFoundError, json.JSONDecodeError):
-                            # Corrupted or incomplete extension, or manifest missing.
-                            pass
+        if not vscode_extensions_dir.exists():
+            return False
+
+        # Iterate through the subdirectories in the extensions directory.
+        extension_dirs = filter(
+            lambda p: p.is_dir() and p.name.startswith(extension_id + "-"),
+            vscode_extensions_dir.iterdir(),
+        )
+        for extension_dir in extension_dirs:
+            # As a more robust check, the manifest file must exist.
+            manifest_path = extension_dir / "package.json"
+            if not manifest_path.exists() or not manifest_path.is_file():
+                continue
+
+            # Finally, the manifest file must be a valid json
+            with open(manifest_path, "r", encoding="utf-8") as f:
+                json.load(f)
+
+            return True
     except Exception:
         pass
 
@@ -72,10 +73,8 @@ def _is_vscode_extension_installed(extension_id: str) -> bool:
 def _is_package_installed(package_name: str) -> bool:
     """
     Checks if a Python package is installed.
-
     Args:
         package_name: The name of the package to check (e.g., "requests", "numpy").
-
     Returns:
         True if the package is installed, False otherwise.
     """
diff --git a/packages/bigframes/tests/unit/session/test_clients.py b/packages/bigframes/tests/unit/session/test_clients.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import os
+import pathlib
+import tempfile
 from typing import Optional
 import unittest.mock as mock
 
@@ -155,6 +157,7 @@ def test_user_agent_not_in_vscode(monkeypatch):
     monkeypatch_client_constructors(monkeypatch)
     provider = create_clients_provider()
     assert_clients_wo_user_agent(provider, "vscode")
+    assert_clients_wo_user_agent(provider, "googlecloudtools.cloudcode")
 
     # We still need to include attribution to bigframes
     assert_clients_w_user_agent(provider, f"bigframes/{bigframes.version.__version__}")
@@ -165,16 +168,48 @@ def test_user_agent_in_vscode(monkeypatch):
     monkeypatch_client_constructors(monkeypatch)
     provider = create_clients_provider()
     assert_clients_w_user_agent(provider, "vscode")
+    assert_clients_wo_user_agent(provider, "googlecloudtools.cloudcode")
 
     # We still need to include attribution to bigframes
     assert_clients_w_user_agent(provider, f"bigframes/{bigframes.version.__version__}")
 
 
+@mock.patch.dict(os.environ, {"VSCODE_PID": "12345"}, clear=True)
+def test_user_agent_in_vscode_w_extension(monkeypatch):
+    monkeypatch_client_constructors(monkeypatch)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        user_home = pathlib.Path(tmpdir)
+        extension_dir = (
+            user_home / ".vscode" / "extensions" / "googlecloudtools.cloudcode-0.12"
+        )
+        extension_config = extension_dir / "package.json"
+
+        # originally extension config does not exist
+        assert not extension_config.exists()
+
+        # simulate extension installation by creating extension config on disk
+        extension_dir.mkdir(parents=True)
+        with open(extension_config, "w") as f:
+            f.write("{}")
+
+        with mock.patch("pathlib.Path.home", return_value=user_home):
+            provider = create_clients_provider()
+            assert_clients_w_user_agent(provider, "vscode")
+            assert_clients_w_user_agent(provider, "googlecloudtools.cloudcode")
+
+            # We still need to include attribution to bigframes
+            assert_clients_w_user_agent(
+                provider, f"bigframes/{bigframes.version.__version__}"
+            )
+
+
 @mock.patch.dict(os.environ, {}, clear=True)
 def test_user_agent_not_in_jupyter(monkeypatch):
     monkeypatch_client_constructors(monkeypatch)
     provider = create_clients_provider()
     assert_clients_wo_user_agent(provider, "jupyter")
+    assert_clients_wo_user_agent(provider, "bigquery_jupyter_plugin")
 
     # We still need to include attribution to bigframes
     assert_clients_w_user_agent(provider, f"bigframes/{bigframes.version.__version__}")
@@ -185,6 +220,37 @@ def test_user_agent_in_jupyter(monkeypatch):
     monkeypatch_client_constructors(monkeypatch)
     provider = create_clients_provider()
     assert_clients_w_user_agent(provider, "jupyter")
+    assert_clients_wo_user_agent(provider, "bigquery_jupyter_plugin")
 
     # We still need to include attribution to bigframes
     assert_clients_w_user_agent(provider, f"bigframes/{bigframes.version.__version__}")
+
+
+@mock.patch.dict(os.environ, {"JPY_PARENT_PID": "12345"}, clear=True)
+def test_user_agent_in_jupyter_with_plugin(monkeypatch):
+    monkeypatch_client_constructors(monkeypatch)
+
+    def custom_import_module_side_effect(name, package=None):
+        if name == "bigquery_jupyter_plugin":
+            return mock.MagicMock()
+        else:
+            import importlib
+
+            return importlib.import_module(name, package)
+
+    assert isinstance(
+        custom_import_module_side_effect("bigquery_jupyter_plugin"), mock.MagicMock
+    )
+    assert custom_import_module_side_effect("bigframes") is bigframes
+
+    with mock.patch(
+        "importlib.import_module", side_effect=custom_import_module_side_effect
+    ):
+        provider = create_clients_provider()
+        assert_clients_w_user_agent(provider, "jupyter")
+        assert_clients_w_user_agent(provider, "bigquery_jupyter_plugin")
+
+        # We still need to include attribution to bigframes
+        assert_clients_w_user_agent(
+            provider, f"bigframes/{bigframes.version.__version__}"
+        )
