From e57565d293e083deeca9485a3f9b341dd39bbf5a Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Tue, 3 Nov 2026 05:41:05 +0000 Subject: [PATCH 1/8] feat(notebooks): add Notebooks API v2 Workbench instance samples and remove legacy v1/ML Engine content Problem: The notebooks/ directory contained an empty notebooks/samples/ skeleton referencing the deprecated Notebooks API v1 Apiary guide, a legacy tutorial targeting the decommissioned Cloud ML Engine API (ml.googleapis.com), and zero Python client library samples for managing Workbench instances with the Notebooks API v2 (google-cloud-notebooks). Solution: - Remove the deprecated notebooks/tutorials/cloud-ml-engine/ tutorial and empty notebooks/samples/ directory. - Update notebooks/README.md to document Workbench instances and link to notebooks/v2/. - Add 24 atomic Python client library samples under notebooks/v2/ using google.cloud.notebooks_v2 (google-cloud-notebooks==1.18.0) covering instance lifecycle, configuration updates, upgradability checks, rollbacks, diagnostics, GPU/custom container/CMEK/Shielded VM creation, Security Blueprint hardening (private VPC, single-user custom service account, disabled root/downloads/terminal), and IAM policy management. - Add unit tests in notebooks/v2/instances_test.py along with requirements.txt, requirements-test.txt, and noxfile_config.py. --- notebooks/README.md | 78 +-- notebooks/samples/README.md | 39 -- notebooks/samples/requirements-test.txt | 1 - notebooks/samples/requirements.txt | 1 - ...ing and prediction with scikit-learn.ipynb | 570 ------------------ notebooks/v2/README.md | 58 ++ notebooks/v2/check_instance_upgradability.py | 83 +++ notebooks/v2/create_instance.py | 89 +++ notebooks/v2/create_instance_cmek.py | 124 ++++ notebooks/v2/create_instance_confidential.py | 99 +++ .../v2/create_instance_custom_container.py | 104 ++++ notebooks/v2/create_instance_gpu.py | 99 +++ notebooks/v2/create_instance_hardened.py | 107 ++++ notebooks/v2/create_instance_private_vpc.py | 117 ++++ notebooks/v2/create_instance_single_user.py | 113 ++++ notebooks/v2/delete_instance.py | 81 +++ notebooks/v2/diagnose_instance.py | 97 +++ notebooks/v2/get_instance.py | 81 +++ notebooks/v2/get_instance_iam_policy.py | 82 +++ notebooks/v2/instances_test.py | 343 +++++++++++ notebooks/v2/list_instances.py | 74 +++ notebooks/v2/noxfile_config.py | 31 + notebooks/v2/requirements-test.txt | 1 + notebooks/v2/requirements.txt | 1 + notebooks/v2/reset_instance.py | 80 +++ notebooks/v2/rollback_instance.py | 101 ++++ notebooks/v2/set_instance_iam_policy.py | 107 ++++ notebooks/v2/start_instance.py | 81 +++ notebooks/v2/stop_instance.py | 80 +++ notebooks/v2/test_instance_iam_permissions.py | 82 +++ notebooks/v2/update_instance_machine_type.py | 96 +++ notebooks/v2/update_instance_metadata.py | 90 +++ .../v2/update_instance_shielded_config.py | 102 ++++ notebooks/v2/upgrade_instance.py | 80 +++ 34 files changed, 2716 insertions(+), 656 deletions(-) delete mode 100644 notebooks/samples/README.md delete mode 100644 notebooks/samples/requirements-test.txt delete mode 100644 notebooks/samples/requirements.txt delete mode 100644 notebooks/tutorials/cloud-ml-engine/Training and prediction with scikit-learn.ipynb create mode 100644 notebooks/v2/README.md create mode 100644 notebooks/v2/check_instance_upgradability.py create mode 100644 notebooks/v2/create_instance.py create mode 100644 notebooks/v2/create_instance_cmek.py create mode 100644 notebooks/v2/create_instance_confidential.py create mode 100644 notebooks/v2/create_instance_custom_container.py create mode 100644 notebooks/v2/create_instance_gpu.py create mode 100644 notebooks/v2/create_instance_hardened.py create mode 100644 notebooks/v2/create_instance_private_vpc.py create mode 100644 notebooks/v2/create_instance_single_user.py create mode 100644 notebooks/v2/delete_instance.py create mode 100644 notebooks/v2/diagnose_instance.py create mode 100644 notebooks/v2/get_instance.py create mode 100644 notebooks/v2/get_instance_iam_policy.py create mode 100644 notebooks/v2/instances_test.py create mode 100644 notebooks/v2/list_instances.py create mode 100644 notebooks/v2/noxfile_config.py create mode 100644 notebooks/v2/requirements-test.txt create mode 100644 notebooks/v2/requirements.txt create mode 100644 notebooks/v2/reset_instance.py create mode 100644 notebooks/v2/rollback_instance.py create mode 100644 notebooks/v2/set_instance_iam_policy.py create mode 100644 notebooks/v2/start_instance.py create mode 100644 notebooks/v2/stop_instance.py create mode 100644 notebooks/v2/test_instance_iam_permissions.py create mode 100644 notebooks/v2/update_instance_machine_type.py create mode 100644 notebooks/v2/update_instance_metadata.py create mode 100644 notebooks/v2/update_instance_shielded_config.py create mode 100644 notebooks/v2/upgrade_instance.py diff --git a/notebooks/README.md b/notebooks/README.md index d00c7fcd337..457942edab4 100644 --- a/notebooks/README.md +++ b/notebooks/README.md @@ -1,82 +1,70 @@ -# AI Platform Notebooks +# Workbench Instances and Notebooks -[AI Platform Notebooks](https://cloud.google.com/ai-platform-notebooks) -is a managed service that offers an integrated and secure JupyterLab -environment for data scientists and machine learning developers to -experiment, develop, and deploy models into production. +**Workbench instances** provide a managed, integrated, and secure JupyterLab environment for data scientists, machine learning developers, and agent builders to experiment, develop, and deploy workloads on Google Cloud. -This directory contains code samples and -tutorials. The tutorials are divided into 2 sections: -- [API code samples](samples) -- [tutorials](tutorials) +This directory contains Python client library code samples and Jupyter notebook tutorials: +- [Notebooks API v2 code samples](v2) +- [Jupyter notebook tutorials](tutorials) -## Notebooks API +## Notebooks API v2 (`notebooks/v2`) -In [this](samples) folder you will find Python code samples to interact -with the -[API](https://cloud.google.com/ai-platform/notebooks/docs/reference/rest) +In the [v2](v2) folder, you will find atomic Python code samples using the `google-cloud-notebooks` (`google.cloud.notebooks_v2`) client library to create, manage, configure, upgrade, and diagnose Workbench instances. -## Tutorials +## Tutorials (`notebooks/tutorials`) -**Run in Notebooks** +### Run in a Workbench Instance -1. Create a Notebook - [instance](https://cloud.google.com/ai-platform/notebooks/docs#how-to) -2. Clone this repository via Jupyter console or using the Git clone - button +1. Create a Workbench instance using the Google Cloud console, the `gcloud workbench instances create` CLI command, or the [create_instance.py](v2/create_instance.py) sample. +2. Clone this repository via the JupyterLab terminal or using the Git clone button: ```bash git clone https://github.com/GoogleCloudPlatform/python-docs-samples.git ``` -3. Open the tutorials folder and open the Jupyter notebooks +3. Open the `notebooks/tutorials` folder and launch the Jupyter notebooks. +### Run Locally -**Run locally** +1. Install JupyterLab or Jupyter Notebook ([installation instructions](https://jupyter.org/install)). +2. Install the dependencies in the [requirements.txt](./requirements.txt) file ([instructions below](#install-the-dependencies)). +3. Register the `google-cloud-bigquery` magic commands ([instructions below](#register-magics-and-configure-matplotlib)). +4. Set `matplotlib` to render inline ([instructions below](#register-magics-and-configure-matplotlib)). -1. Install Jupyter notebooks - ([instructions](https://jupyter.org/install)) -1. Install the dependencies in the [requirements.txt](./requirements.txt) file ([instructions below](#install-the-dependencies)) -1. Registered the `google-cloud-bigquery` magic commands ([instructions below](#register-magics-and-configure-matplotlib)) -1. Set `matplotlib` to render inline ([instructions below](#register-magics-and-configure-matplotlib)) - -### Install the dependencies +#### Install the dependencies Install the dependencies with the following command: - pip install --upgrade -r requirements.txt +```bash +pip install --upgrade -r requirements.txt +``` + +#### Register magics and configure matplotlib -### Register magics and configure matplotlib +You can either perform these setup steps in a single notebook, or add the steps to your IPython configuration file to apply to all notebooks. -You can either perform these set up steps in a single notebook, or add the -steps to your IPython configuration file to apply to all notebooks. +##### Perform setup steps within a notebook -### Perform set up steps within a notebook +Run the following commands in your notebook to register the BigQuery magic commands and set `matplotlib` to render inline: -To perform the set up steps for a single notebook, run the following commands -in your notebook to register the BigQuery magic commands and set `matplotlib` -to render inline: ```python %load_ext google.cloud.bigquery %matplotlib inline ``` -### Perform set up steps in your IPython configuration file +##### Perform setup steps in your IPython configuration file + +Add the following code to your `ipython_config.py` file to register the BigQuery magic commands and set `matplotlib` to render inline across all notebooks: -To perform the set up steps implicitly for all of your notebooks, add the -following code to your `ipython_config.py` file to register the BigQuery magic -commands and set `matplotlib` to render inline: ```python c = get_config() # Register magic commands c.InteractiveShellApp.extensions = [ - 'google.cloud.bigquery', + "google.cloud.bigquery", ] # Enable matplotlib renderings to render inline in the notebook. -c.InteractiveShellApp.matplotlib = 'inline' +c.InteractiveShellApp.matplotlib = "inline" ``` -See -[IPython documentation](https://ipython.readthedocs.io/en/stable/config/intro.html) -for more information about IPython configuration. + +See the [IPython configuration documentation](https://ipython.readthedocs.io/en/stable/config/intro.html) for more information. diff --git a/notebooks/samples/README.md b/notebooks/samples/README.md deleted file mode 100644 index 7a92a46d3d0..00000000000 --- a/notebooks/samples/README.md +++ /dev/null @@ -1,39 +0,0 @@ -# API code samples - -Sample command-line programs for interacting with the Notebooks API. - -## Prerequisites to run locally: - -* [Python 3](https://www.python.org/downloads/) -* [pip](https://pypi.python.org/pypi/pip) - -Go to the [Google Cloud Console](https://console.cloud.google.com). - -Under API Manager, search for the Google Cloud Notebooks API and enable -it. You can also enable API via Command Line: - -```bash -gcloud services enable notebooks.googleapis.com -``` - -## Set Up Your Local Dev Environment - -To install, run the following commands. If you want to use -[virtualenv](https://virtualenv.readthedocs.org/en/latest/) -(recommended), run the commands within a virtualenv. - -```bash -pip install -r requirements.txt -``` - -**Note:** While this sample demonstrates interacting with Notebooks via -the API, the functionality demonstrated here could also be accomplished -using the Cloud Console or the -[gcloud CLI](https://cloud.google.com/sdk/gcloud/reference/beta/notebooks) - -## Authentication - -Please see the -[Google cloud authentication guide](https://cloud.google.com/docs/authentication/). -The recommended approach to running these samples is a Service Account -with a JSON key. diff --git a/notebooks/samples/requirements-test.txt b/notebooks/samples/requirements-test.txt deleted file mode 100644 index c9e154ba440..00000000000 --- a/notebooks/samples/requirements-test.txt +++ /dev/null @@ -1 +0,0 @@ -pytest==9.0.3; python_version >= "3.10" diff --git a/notebooks/samples/requirements.txt b/notebooks/samples/requirements.txt deleted file mode 100644 index 76dc107a22b..00000000000 --- a/notebooks/samples/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -google-api-python-client==2.131.0 diff --git a/notebooks/tutorials/cloud-ml-engine/Training and prediction with scikit-learn.ipynb b/notebooks/tutorials/cloud-ml-engine/Training and prediction with scikit-learn.ipynb deleted file mode 100644 index d805d0cb1fe..00000000000 --- a/notebooks/tutorials/cloud-ml-engine/Training and prediction with scikit-learn.ipynb +++ /dev/null @@ -1,570 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Training and prediction with scikit-learn\n", - "\n", - "This notebook demonstrates how to use AI Platform to train a simple classification model using `scikit-learn`, and then deploy the model to get predictions.\n", - "\n", - "You train the model to predict a person's income level based on the [Census Income data set](https://archive.ics.uci.edu/ml/datasets/Census+Income).\n", - "\n", - "Before you jump in, let’s cover some of the different tools you’ll be using:\n", - "\n", - "+ [AI Platform](https://cloud.google.com/ml-engine/) is a managed service that enables you to easily build machine learning models that work on any type of data, of any size.\n", - "\n", - "+ [Cloud Storage](https://cloud.google.com/storage/) is a unified object storage for developers and enterprises, from live data serving to data analytics/ML to data archiving.\n", - "\n", - "+ [Cloud SDK](https://cloud.google.com/sdk/) is a command line tool which allows you to interact with Google Cloud products. This notebook introduces several `gcloud` and `gsutil` commands, which are part of the Cloud SDK. Note that shell commands in a notebook must be prepended with a `!`." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Set up your environment\n", - "\n", - "### Enable the required APIs\n", - "\n", - "In order to use AI Platform, confirm that the required APIs are enabled:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!gcloud services enable ml.googleapis.com\n", - "!gcloud services enable compute.googleapis.com" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Create a storage bucket\n", - "Buckets are the basic containers that hold your data. Everything that you store in Cloud Storage must be contained in a bucket. You can use buckets to organize your data and control access to your data.\n", - "\n", - "Start by defining a globally unique name.\n", - "\n", - "For more information about naming buckets, see [Bucket name requirements](https://cloud.google.com/storage/docs/naming#requirements)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "BUCKET_NAME = \"your-new-bucket\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the examples below, the `BUCKET_NAME` variable is referenced in the commands using `$`.\n", - "\n", - "Create the new bucket with the `gsutil mb` command:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!gcloud storage buckets create gs://$BUCKET_NAME" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### About the data\n", - "\n", - "The [Census Income Data Set](https://archive.ics.uci.edu/ml/datasets/Census+Income) that this sample\n", - "uses for training is provided by the [UC Irvine Machine Learning\n", - "Repository](https://archive.ics.uci.edu/ml/datasets/).\n", - "\n", - "Census data courtesy of: Lichman, M. (2013). UCI Machine Learning Repository http://archive.ics.uci.edu/ml. Irvine, CA: University of California, School of Information and Computer Science. This dataset is publicly available for anyone to use under the following terms provided by the Dataset Source - http://archive.ics.uci.edu/ml - and is provided \"AS IS\" without any warranty, express or implied, from Google. Google disclaims all liability for any damages, direct or indirect, resulting from the use of the dataset.\n", - "\n", - "The data used in this tutorial is located in a public Cloud Storage bucket:\n", - "\n", - " gs://cloud-samples-data/ml-engine/sklearn/census_data/ \n", - "\n", - "The training file is `adult.data` ([download](https://storage.googleapis.com/cloud-samples-data/ml-engine/sklearn/census_data/adult.data)) and the evaluation file is `adult.test` ([download](https://storage.googleapis.com/cloud-samples-data/ml-engine/sklearn/census_data/adult.test)). The evaluation file is not used in this tutorial." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Create training application package\n", - "\n", - "The easiest (and recommended) way to create a training application package is to use `gcloud` to package and upload the application when you submit your training job. This method allows you to create a very simple file structure with only two files. For this tutorial, the file structure of your training application package should appear similar to the following:\n", - "\n", - " census_training/\n", - " __init__.py\n", - " train.py" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create a directory locally:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!mkdir census_training" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create a blank file named `__init__.py`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!touch ./census_training/__init__.py" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Save training code in one Python file in the `census_training` directory. The following cell writes a training file to the `census_training` directory. The training file performs the following operations:\n", - "+ Loads the data into a pandas `DataFrame` that can be used by `scikit-learn`\n", - "+ Fits the model is against the training data\n", - "+ Exports the model with the [Python `pickle` library](https://docs.python.org/3/library/pickle.html)\n", - "\n", - "The following model training code is not executed within this notebook. Instead, it is saved to a Python file and packaged as a Python module that runs on AI Platform after you submit the training job." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%writefile ./census_training/train.py\n", - "import argparse\n", - "import pickle\n", - "import pandas as pd\n", - "\n", - "from google.cloud import storage\n", - "\n", - "from sklearn.ensemble import RandomForestClassifier\n", - "from sklearn.feature_selection import SelectKBest\n", - "from sklearn.pipeline import FeatureUnion\n", - "from sklearn.pipeline import Pipeline\n", - "from sklearn.preprocessing import LabelBinarizer\n", - "\n", - "parser = argparse.ArgumentParser()\n", - "parser.add_argument(\"--bucket-name\", help=\"The bucket name\", required=True)\n", - "\n", - "arguments, unknown = parser.parse_known_args()\n", - "bucket_name = arguments.bucket_name\n", - "\n", - "# Define the format of your input data, including unused columns.\n", - "# These are the columns from the census data files.\n", - "COLUMNS = (\n", - " 'age',\n", - " 'workclass',\n", - " 'fnlwgt',\n", - " 'education',\n", - " 'education-num',\n", - " 'marital-status',\n", - " 'occupation',\n", - " 'relationship',\n", - " 'race',\n", - " 'sex',\n", - " 'capital-gain',\n", - " 'capital-loss',\n", - " 'hours-per-week',\n", - " 'native-country',\n", - " 'income-level'\n", - ")\n", - "\n", - "# Categorical columns are columns that need to be turned into a numerical value\n", - "# to be used by scikit-learn\n", - "CATEGORICAL_COLUMNS = (\n", - " 'workclass',\n", - " 'education',\n", - " 'marital-status',\n", - " 'occupation',\n", - " 'relationship',\n", - " 'race',\n", - " 'sex',\n", - " 'native-country'\n", - ")\n", - "\n", - "# Create a Cloud Storage client to download the census data\n", - "storage_client = storage.Client()\n", - "\n", - "# Download the data\n", - "public_bucket = storage_client.bucket('cloud-samples-data')\n", - "blob = public_bucket.blob('ml-engine/sklearn/census_data/adult.data')\n", - "blob.download_to_filename('adult.data')\n", - "\n", - "# Load the training census dataset\n", - "with open(\"./adult.data\", \"r\") as train_data:\n", - " raw_training_data = pd.read_csv(train_data, header=None, names=COLUMNS)\n", - " # Removing the whitespaces in categorical features\n", - " for col in CATEGORICAL_COLUMNS:\n", - " raw_training_data[col] = raw_training_data[col].apply(lambda x: str(x).strip())\n", - "\n", - "# Remove the column we are trying to predict ('income-level') from our features\n", - "# list and convert the DataFrame to a lists of lists\n", - "train_features = raw_training_data.drop(\"income-level\", axis=1).values.tolist()\n", - "# Create our training labels list, convert the DataFrame to a lists of lists\n", - "train_labels = (raw_training_data[\"income-level\"] == \" >50K\").values.tolist()\n", - "\n", - "# Since the census data set has categorical features, we need to convert\n", - "# them to numerical values. We'll use a list of pipelines to convert each\n", - "# categorical column and then use FeatureUnion to combine them before calling\n", - "# the RandomForestClassifier.\n", - "categorical_pipelines = []\n", - "\n", - "# Each categorical column needs to be extracted individually and converted to a\n", - "# numerical value. To do this, each categorical column will use a pipeline that\n", - "# extracts one feature column via SelectKBest(k=1) and a LabelBinarizer() to\n", - "# convert the categorical value to a numerical one. A scores array (created\n", - "# below) will select and extract the feature column. The scores array is\n", - "# created by iterating over the columns and checking if it is a\n", - "# categorical column.\n", - "for i, col in enumerate(COLUMNS[:-1]):\n", - " if col in CATEGORICAL_COLUMNS:\n", - " # Create a scores array to get the individual categorical column.\n", - " # Example:\n", - " # data = [\n", - " # 39, 'State-gov', 77516, 'Bachelors', 13, 'Never-married',\n", - " # 'Adm-clerical', 'Not-in-family', 'White', 'Male', 2174, 0,\n", - " # 40, 'United-States'\n", - " # ]\n", - " # scores = [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n", - " #\n", - " # Returns: [['State-gov']]\n", - " # Build the scores array\n", - " scores = [0] * len(COLUMNS[:-1])\n", - " # This column is the categorical column we want to extract.\n", - " scores[i] = 1\n", - " skb = SelectKBest(k=1)\n", - " skb.scores_ = scores\n", - " # Convert the categorical column to a numerical value\n", - " lbn = LabelBinarizer()\n", - " r = skb.transform(train_features)\n", - " lbn.fit(r)\n", - " # Create the pipeline to extract the categorical feature\n", - " categorical_pipelines.append(\n", - " (\n", - " 'categorical-{}'.format(i), \n", - " Pipeline([\n", - " ('SKB-{}'.format(i), skb),\n", - " ('LBN-{}'.format(i), lbn)])\n", - " )\n", - " )\n", - "\n", - "# Create pipeline to extract the numerical features\n", - "skb = SelectKBest(k=6)\n", - "# From COLUMNS use the features that are numerical\n", - "skb.scores_ = [1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0]\n", - "categorical_pipelines.append((\"numerical\", skb))\n", - "\n", - "# Combine all the features using FeatureUnion\n", - "preprocess = FeatureUnion(categorical_pipelines)\n", - "\n", - "# Create the classifier\n", - "classifier = RandomForestClassifier()\n", - "\n", - "# Transform the features and fit them to the classifier\n", - "classifier.fit(preprocess.transform(train_features), train_labels)\n", - "\n", - "# Create the overall model as a single pipeline\n", - "pipeline = Pipeline([(\"union\", preprocess), (\"classifier\", classifier)])\n", - "\n", - "# Create the model file\n", - "# It is required to name the model file \"model.pkl\" if you are using pickle\n", - "model_filename = \"model.pkl\"\n", - "with open(model_filename, \"wb\") as model_file:\n", - " pickle.dump(pipeline, model_file)\n", - "\n", - "# Upload the model to Cloud Storage\n", - "bucket = storage_client.bucket(bucket_name)\n", - "blob = bucket.blob(model_filename)\n", - "blob.upload_from_filename(model_filename)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Submit the training job\n", - "\n", - "In this section, you use [`gcloud ai-platform jobs submit training`](https://cloud.google.com/sdk/gcloud/reference/ai-platform/jobs/submit/training) to submit your training job. The `--` argument passed to the command is a separator; anything after the separator will be passed to the Python code as input arguments.\n", - "\n", - "For more information about the arguments preceeding the separator, run the following:\n", - "\n", - " gcloud ai-platform jobs submit training --help\n", - "\n", - "The argument given to the python script is `--bucket-name`. The `--bucket-name` argument is used to specify the name of the bucket to save the model file." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "import time\n", - "\n", - "# Define a timestamped job name\n", - "JOB_NAME = \"census_training_{}\".format(int(time.time()))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "# Submit the training job:\n", - "!gcloud ai-platform jobs submit training $JOB_NAME \\\n", - " --job-dir gs://$BUCKET_NAME/census_job_dir \\\n", - " --package-path ./census_training \\\n", - " --module-name census_training.train \\\n", - " --region us-central1 \\\n", - " --runtime-version=1.12 \\\n", - " --python-version=3.5 \\\n", - " --scale-tier BASIC \\\n", - " --stream-logs \\\n", - " -- \\\n", - " --bucket-name $BUCKET_NAME" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Verify model file in Cloud Storage\n", - "\n", - "View the contents of the destination model directory to verify that your model file has been uploaded to Cloud Storage.\n", - "\n", - "Note: The model can take a few minutes to train and show up in Cloud Storage." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!gcloud storage ls gs://$BUCKET_NAME/" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Serve the model" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Once the model is successfully created and trained, you can serve it. A model can have different versions. In order to serve the model, create a model and version in AI Platform.\n", - "\n", - "Define the model and version names:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "MODEL_NAME = \"CensusPredictor\"\n", - "VERSION_NAME = \"census_predictor_{}\".format(int(time.time()))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create the model in AI Platform:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!gcloud ai-platform models create $MODEL_NAME --regions us-central1" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Create a version that points to your model file in Cloud Storage:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!gcloud ai-platform versions create $VERSION_NAME \\\n", - " --model=$MODEL_NAME \\\n", - " --framework=scikit-learn \\\n", - " --origin=gs://$BUCKET_NAME/ \\\n", - " --python-version=3.5 \\\n", - " --runtime-version=1.12" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Make predictions\n", - "\n", - "### Format data for prediction\n", - "\n", - "Before you send an online prediction request, you must format your test data to prepare it for use by the AI Platform prediction service. Make sure that the format of your input instances matches what your model expects.\n", - "\n", - "Create an `input.json` file with each input instance on a separate line. The following example uses ten data instances. Note that the format of input instances needs to match what your model expects. In this example, the Census model requires 14 features, so your input must be a matrix of shape (`num_instances, 14`)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Define a name for the input file\n", - "INPUT_FILE = \"./census_training/input.json\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "%%writefile $INPUT_FILE\n", - "[25, \"Private\", 226802, \"11th\", 7, \"Never-married\", \"Machine-op-inspct\", \"Own-child\", \"Black\", \"Male\", 0, 0, 40, \"United-States\"]\n", - "[38, \"Private\", 89814, \"HS-grad\", 9, \"Married-civ-spouse\", \"Farming-fishing\", \"Husband\", \"White\", \"Male\", 0, 0, 50, \"United-States\"]\n", - "[28, \"Local-gov\", 336951, \"Assoc-acdm\", 12, \"Married-civ-spouse\", \"Protective-serv\", \"Husband\", \"White\", \"Male\", 0, 0, 40, \"United-States\"]\n", - "[44, \"Private\", 160323, \"Some-college\", 10, \"Married-civ-spouse\", \"Machine-op-inspct\", \"Husband\", \"Black\", \"Male\", 7688, 0, 40, \"United-States\"]\n", - "[18, \"?\", 103497, \"Some-college\", 10, \"Never-married\", \"?\", \"Own-child\", \"White\", \"Female\", 0, 0, 30, \"United-States\"]\n", - "[34, \"Private\", 198693, \"10th\", 6, \"Never-married\", \"Other-service\", \"Not-in-family\", \"White\", \"Male\", 0, 0, 30, \"United-States\"]\n", - "[29, \"?\", 227026, \"HS-grad\", 9, \"Never-married\", \"?\", \"Unmarried\", \"Black\", \"Male\", 0, 0, 40, \"United-States\"]\n", - "[63, \"Self-emp-not-inc\", 104626, \"Prof-school\", 15, \"Married-civ-spouse\", \"Prof-specialty\", \"Husband\", \"White\", \"Male\", 3103, 0, 32, \"United-States\"]\n", - "[24, \"Private\", 369667, \"Some-college\", 10, \"Never-married\", \"Other-service\", \"Unmarried\", \"White\", \"Female\", 0, 0, 40, \"United-States\"]\n", - "[55, \"Private\", 104996, \"7th-8th\", 4, \"Married-civ-spouse\", \"Craft-repair\", \"Husband\", \"White\", \"Male\", 0, 0, 10, \"United-States\"]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Send the online prediction request\n", - "\n", - "The prediction results return `True` if the person's income is predicted to be greater than $50,000 per year, and `False` otherwise. The output of the command below may appear similar to the following:\n", - "\n", - " [False, False, False, True, False, False, False, False, False, False]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "!gcloud ai-platform predict --model $MODEL_NAME --version \\\n", - " $VERSION_NAME --json-instances $INPUT_FILE" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Clean up" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To delete all resources you created in this tutorial, run the following commands:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "# Delete the model version\n", - "!gcloud ai-platform versions delete $VERSION_NAME --model=$MODEL_NAME --quiet\n", - "\n", - "# Delete the model\n", - "!gcloud ai-platform models delete $MODEL_NAME --quiet\n", - "\n", - "# Delete the bucket and contents\n", - "!gcloud storage rm --recursive gs://$BUCKET_NAME\n", - "\n", - "# Delete the local files created by the tutorial\n", - "!rm -rf census_training" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.4" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/notebooks/v2/README.md b/notebooks/v2/README.md new file mode 100644 index 00000000000..fb44cca16ba --- /dev/null +++ b/notebooks/v2/README.md @@ -0,0 +1,58 @@ +# Workbench Instance Samples (Notebooks API v2) + +This directory contains atomic Python code samples demonstrating how to manage **Workbench instances** using the `google-cloud-notebooks` (`google.cloud.notebooks_v2`) client library. + +## Prerequisites + +1. Enable the **Notebooks API** (`notebooks.googleapis.com`) in your Google Cloud project: + ```bash + gcloud services enable notebooks.googleapis.com + ``` +2. Configure Application Default Credentials (ADC): + ```bash + gcloud auth application-default login + ``` +3. Create a virtual environment and install dependencies: + ```bash + python3 -m venv .venv + .venv/bin/pip install -r requirements.txt + ``` + +## Samples + +### Core Instance Lifecycle & Read Operations +- [`list_instances.py`](list_instances.py) — Lists Workbench instances in a zone (`notebooks_v2_list_instances`). +- [`get_instance.py`](get_instance.py) — Gets metadata and health state for a Workbench instance (`notebooks_v2_get_instance`). +- [`create_instance.py`](create_instance.py) — Creates a new Workbench instance with default VM settings (`notebooks_v2_create_instance`). +- [`stop_instance.py`](stop_instance.py) — Stops a running Workbench instance (`notebooks_v2_stop_instance`). +- [`start_instance.py`](start_instance.py) — Starts a stopped Workbench instance (`notebooks_v2_start_instance`). +- [`reset_instance.py`](reset_instance.py) — Resets (reboots) a Workbench instance (`notebooks_v2_reset_instance`). +- [`delete_instance.py`](delete_instance.py) — Deletes a Workbench instance (`notebooks_v2_delete_instance`). + +### Instance Configuration, Maintenance & Diagnostics +- [`update_instance_machine_type.py`](update_instance_machine_type.py) — Updates the Compute Engine machine type of a stopped Workbench instance (`notebooks_v2_update_instance_machine_type`). +- [`update_instance_metadata.py`](update_instance_metadata.py) — Updates custom metadata keys on a Workbench instance (`notebooks_v2_update_instance_metadata`). +- [`update_instance_shielded_config.py`](update_instance_shielded_config.py) — Updates the Shielded VM configuration on a stopped Workbench instance (`notebooks_v2_update_instance_shielded_config`). +- [`check_instance_upgradability.py`](check_instance_upgradability.py) — Checks whether a Workbench instance can be upgraded (`notebooks_v2_check_instance_upgradability`). +- [`upgrade_instance.py`](upgrade_instance.py) — Upgrades a Workbench instance to the latest environment version (`notebooks_v2_upgrade_instance`). +- [`rollback_instance.py`](rollback_instance.py) — Rolls back a Workbench instance to a previous snapshot and revision (`notebooks_v2_rollback_instance`). +- [`diagnose_instance.py`](diagnose_instance.py) — Runs a diagnostic bundle export to Cloud Storage for a Workbench instance (`notebooks_v2_diagnose_instance`). + +### Specialized Creation, Security Blueprint Hardening & Access Control +- [`create_instance_gpu.py`](create_instance_gpu.py) — Creates a GPU-accelerated Workbench instance (`notebooks_v2_create_instance_gpu`). +- [`create_instance_custom_container.py`](create_instance_custom_container.py) — Creates a Workbench instance from a custom container image (`notebooks_v2_create_instance_custom_container`). +- [`create_instance_cmek.py`](create_instance_cmek.py) — Creates a Workbench instance encrypted with a customer-managed encryption key (`notebooks_v2_create_instance_cmek`). +- [`create_instance_confidential.py`](create_instance_confidential.py) — Creates a Shielded VM Workbench instance (`notebooks_v2_create_instance_confidential`). +- [`create_instance_private_vpc.py`](create_instance_private_vpc.py) — Creates a private-IP Workbench instance in a custom VPC network and subnet (`notebooks_v2_create_instance_private_vpc`). +- [`create_instance_single_user.py`](create_instance_single_user.py) — Creates a single-user Workbench instance with a custom least-privilege service account (`notebooks_v2_create_instance_single_user`). +- [`create_instance_hardened.py`](create_instance_hardened.py) — Creates a security-hardened Workbench instance disabling root, downloads, terminal, and public IP (`notebooks_v2_create_instance_hardened`). +- [`get_instance_iam_policy.py`](get_instance_iam_policy.py) — Gets the IAM policy for a Workbench instance (`notebooks_v2_get_instance_iam_policy`). +- [`set_instance_iam_policy.py`](set_instance_iam_policy.py) — Sets an IAM policy binding on a Workbench instance (`notebooks_v2_set_instance_iam_policy`). +- [`test_instance_iam_permissions.py`](test_instance_iam_permissions.py) — Tests IAM permissions on a Workbench instance (`notebooks_v2_test_instance_iam_permissions`). + +## Running Tests + +```bash +.venv/bin/pip install -r requirements-test.txt +.venv/bin/pytest instances_test.py +``` diff --git a/notebooks/v2/check_instance_upgradability.py b/notebooks/v2/check_instance_upgradability.py new file mode 100644 index 00000000000..b41b19aac5d --- /dev/null +++ b/notebooks/v2/check_instance_upgradability.py @@ -0,0 +1,83 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_check_instance_upgradability] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def check_instance_upgradability( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Checks whether a Workbench instance can be upgraded. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to check. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.CheckInstanceUpgradabilityRequest( + notebook_instance=name, + ) + + try: + response = client.check_instance_upgradability(request=request) + print(f"Resource ID: {name}") + print(f"Upgrade Image: {response.upgrade_image}") + print(f"Upgrade Info: {response.upgrade_info}") + print(f"Upgrade Version: {response.upgrade_version}") + print(f"Upgradeable: {response.upgradeable}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_check_instance_upgradability] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Check whether a Workbench instance can be upgraded." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + check_instance_upgradability( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/create_instance.py b/notebooks/v2/create_instance.py new file mode 100644 index 00000000000..0218e742970 --- /dev/null +++ b/notebooks/v2/create_instance.py @@ -0,0 +1,89 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_create_instance] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Creates a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + gce_setup = notebooks_v2.GceSetup(machine_type="e2-standard-4") + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") + print(f"Machine Type: {response.gce_setup.machine_type}") + print(f"Proxy URI: {response.proxy_uri}") + print(f"State: {response.state.name}") + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID for the new Workbench instance.", + ) + args = parser.parse_args() + create_instance( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/create_instance_cmek.py b/notebooks/v2/create_instance_cmek.py new file mode 100644 index 00000000000..033c7bf04fe --- /dev/null +++ b/notebooks/v2/create_instance_cmek.py @@ -0,0 +1,124 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_create_instance_cmek] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_cmek( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-cmek-instance", + key_ring_id: str = "example-key-ring", + crypto_key_id: str = "example-key", +) -> None: + """Creates a CMEK-encrypted Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + key_ring_id: The Cloud KMS key ring ID in the same region. + crypto_key_id: The Cloud KMS crypto key ID used for disk encryption. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + region_id = "-".join(location_id.split("-")[:-1]) + kms_key = ( + f"projects/{project_id}/locations/{region_id}/" + f"keyRings/{key_ring_id}/cryptoKeys/{crypto_key_id}" + ) + boot_disk = notebooks_v2.BootDisk( + disk_size_gb=150, + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=kms_key, + ) + data_disk = notebooks_v2.DataDisk( + disk_size_gb=100, + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=kms_key, + ) + gce_setup = notebooks_v2.GceSetup( + machine_type="e2-standard-4", + boot_disk=boot_disk, + data_disks=[data_disk], + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + encryption = response.gce_setup.boot_disk.disk_encryption.name + print(f"Resource ID: {response.name}") + print(f"Boot Disk Encryption: {encryption}") + print(f"KMS Key: {response.gce_setup.boot_disk.kms_key}") + print(f"State: {response.state.name}") + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_cmek] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create a CMEK-encrypted Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-cmek-instance", + help="The ID for the new Workbench instance.", + ) + parser.add_argument( + "--key-ring-id", + default="example-key-ring", + help="The Cloud KMS key ring ID.", + ) + parser.add_argument( + "--crypto-key-id", + default="example-key", + help="The Cloud KMS crypto key ID.", + ) + args = parser.parse_args() + create_instance_cmek( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + key_ring_id=args.key_ring_id, + crypto_key_id=args.crypto_key_id, + ) diff --git a/notebooks/v2/create_instance_confidential.py b/notebooks/v2/create_instance_confidential.py new file mode 100644 index 00000000000..01eaa6aedfa --- /dev/null +++ b/notebooks/v2/create_instance_confidential.py @@ -0,0 +1,99 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_create_instance_confidential] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_confidential( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-confidential-instance", +) -> None: + """Creates a Shielded VM Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + shielded_config = notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ) + gce_setup = notebooks_v2.GceSetup( + machine_type="n2d-standard-4", + shielded_instance_config=shielded_config, + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + secure_boot = ( + response.gce_setup.shielded_instance_config.enable_secure_boot + ) + print(f"Resource ID: {response.name}") + print(f"Machine Type: {response.gce_setup.machine_type}") + print(f"Secure Boot Enabled: {secure_boot}") + print(f"State: {response.state.name}") + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_confidential] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create a Shielded and Confidential VM Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-confidential-instance", + help="The ID for the new Workbench instance.", + ) + args = parser.parse_args() + create_instance_confidential( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/create_instance_custom_container.py b/notebooks/v2/create_instance_custom_container.py new file mode 100644 index 00000000000..45ef91e9383 --- /dev/null +++ b/notebooks/v2/create_instance_custom_container.py @@ -0,0 +1,104 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_create_instance_custom_container] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_custom_container( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-container-instance", + repository: str = "gcr.io/deeplearning-platform-release/base-cpu", +) -> None: + """Creates a Workbench instance from a custom container. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + repository: The container image repository URI. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + container_image = notebooks_v2.ContainerImage( + repository=repository, + tag="latest", + ) + gce_setup = notebooks_v2.GceSetup( + machine_type="e2-standard-4", + container_image=container_image, + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print( + f"Container Repository: {response.gce_setup.container_image.repository}" + ) + print(f"State: {response.state.name}") + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_custom_container] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create a Workbench instance using a custom container image." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-container-instance", + help="The ID for the new Workbench instance.", + ) + parser.add_argument( + "--repository", + default="gcr.io/deeplearning-platform-release/base-cpu", + help="The container image repository URI.", + ) + args = parser.parse_args() + create_instance_custom_container( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + repository=args.repository, + ) diff --git a/notebooks/v2/create_instance_gpu.py b/notebooks/v2/create_instance_gpu.py new file mode 100644 index 00000000000..9ee3e851274 --- /dev/null +++ b/notebooks/v2/create_instance_gpu.py @@ -0,0 +1,99 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_create_instance_gpu] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_gpu( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-gpu-instance", +) -> None: + """Creates a GPU-accelerated Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new GPU instance. + instance_id: The ID to assign to the new Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + accelerator_config = notebooks_v2.AcceleratorConfig( + type_=notebooks_v2.AcceleratorConfig.AcceleratorType.NVIDIA_TESLA_T4, + core_count=1, + ) + gpu_driver_config = notebooks_v2.GPUDriverConfig( + enable_gpu_driver=True, + ) + gce_setup = notebooks_v2.GceSetup( + machine_type="n1-standard-4", + accelerator_configs=[accelerator_config], + gpu_driver_config=gpu_driver_config, + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") + print(f"Machine Type: {response.gce_setup.machine_type}") + print(f"State: {response.state.name}") + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_gpu] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create a GPU-accelerated Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-gpu-instance", + help="The ID for the new Workbench instance.", + ) + args = parser.parse_args() + create_instance_gpu( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/create_instance_hardened.py b/notebooks/v2/create_instance_hardened.py new file mode 100644 index 00000000000..8c185c0364f --- /dev/null +++ b/notebooks/v2/create_instance_hardened.py @@ -0,0 +1,107 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_create_instance_hardened] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_hardened( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-hardened-instance", +) -> None: + """Creates a security-hardened Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + shielded_config = notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ) + hardened_metadata = { + "notebook-disable-root": "true", + "notebook-disable-downloads": "true", + "notebook-disable-terminal": "true", + "serial-port-logging-enable": "false", + } + gce_setup = notebooks_v2.GceSetup( + machine_type="e2-standard-4", + disable_public_ip=True, + shielded_instance_config=shielded_config, + metadata=hardened_metadata, + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + secure_boot = ( + response.gce_setup.shielded_instance_config.enable_secure_boot + ) + print(f"Resource ID: {response.name}") + print(f"Disable Public IP: {response.gce_setup.disable_public_ip}") + print(f"Secure Boot Enabled: {secure_boot}") + print(f"State: {response.state.name}") + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_hardened] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create a security-hardened Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-hardened-instance", + help="The ID for the new Workbench instance.", + ) + args = parser.parse_args() + create_instance_hardened( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/create_instance_private_vpc.py b/notebooks/v2/create_instance_private_vpc.py new file mode 100644 index 00000000000..1130196f2e0 --- /dev/null +++ b/notebooks/v2/create_instance_private_vpc.py @@ -0,0 +1,117 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_create_instance_private_vpc] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_private_vpc( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-private-instance", + network_id: str = "example-network", + subnet_id: str = "example-subnet", +) -> None: + """Creates a private-IP Workbench instance in a custom VPC. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + network_id: The ID of the VPC network in the project. + subnet_id: The ID of the subnetwork in the same region. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + region_id = "-".join(location_id.split("-")[:-1]) + network = f"projects/{project_id}/global/networks/{network_id}" + subnet = ( + f"projects/{project_id}/regions/{region_id}/subnetworks/{subnet_id}" + ) + network_interface = notebooks_v2.NetworkInterface( + network=network, + subnet=subnet, + nic_type=notebooks_v2.NetworkInterface.NicType.GVNIC, + ) + gce_setup = notebooks_v2.GceSetup( + machine_type="e2-standard-4", + disable_public_ip=True, + network_interfaces=[network_interface], + ) + instance = notebooks_v2.Instance(gce_setup=gce_setup) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"Disable Public IP: {response.gce_setup.disable_public_ip}") + print(f"State: {response.state.name}") + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_private_vpc] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create a private-IP Workbench instance in a custom VPC." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-private-instance", + help="The ID for the new Workbench instance.", + ) + parser.add_argument( + "--network-id", + default="example-network", + help="The VPC network ID.", + ) + parser.add_argument( + "--subnet-id", + default="example-subnet", + help="The subnetwork ID.", + ) + args = parser.parse_args() + create_instance_private_vpc( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + network_id=args.network_id, + subnet_id=args.subnet_id, + ) diff --git a/notebooks/v2/create_instance_single_user.py b/notebooks/v2/create_instance_single_user.py new file mode 100644 index 00000000000..eb5490d67e1 --- /dev/null +++ b/notebooks/v2/create_instance_single_user.py @@ -0,0 +1,113 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_create_instance_single_user] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def create_instance_single_user( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-single-user-instance", + instance_owner: str = "example-user@example.com", + service_account_id: str = "example-sa", +) -> None: + """Creates a single-user Workbench instance with a custom service account. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone for the new instance. + instance_id: The ID to assign to the new Workbench instance. + instance_owner: The email address of the single permitted user. + service_account_id: The ID of the custom service account in the project. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + service_account_email = ( + f"{service_account_id}@{project_id}.iam.gserviceaccount.com" + ) + service_account = notebooks_v2.ServiceAccount(email=service_account_email) + gce_setup = notebooks_v2.GceSetup( + machine_type="e2-standard-4", + service_accounts=[service_account], + ) + instance = notebooks_v2.Instance( + gce_setup=gce_setup, + instance_owners=[instance_owner], + ) + request = notebooks_v2.CreateInstanceRequest( + parent=parent, + instance_id=instance_id, + instance=instance, + ) + + try: + operation = client.create_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"Creator: {response.creator}") + print(f"State: {response.state.name}") + except exceptions.AlreadyExists as err: + print( + f"error: Instance '{instance_id}' already exists in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_create_instance_single_user] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create a single-user Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-single-user-instance", + help="The ID for the new Workbench instance.", + ) + parser.add_argument( + "--instance-owner", + default="example-user@example.com", + help="The email address of the instance owner.", + ) + parser.add_argument( + "--service-account-id", + default="example-sa", + help="The ID of the custom VM service account.", + ) + args = parser.parse_args() + create_instance_single_user( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + instance_owner=args.instance_owner, + service_account_id=args.service_account_id, + ) diff --git a/notebooks/v2/delete_instance.py b/notebooks/v2/delete_instance.py new file mode 100644 index 00000000000..95e33b95c25 --- /dev/null +++ b/notebooks/v2/delete_instance.py @@ -0,0 +1,81 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_delete_instance] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def delete_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Deletes a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to delete. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.DeleteInstanceRequest(name=name) + + try: + operation = client.delete_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + operation.result() + print(f"Resource ID: {name}") + print(f"Location: {location_id}") + print(f"Status: Deleted") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_delete_instance] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Delete a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + delete_instance( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/diagnose_instance.py b/notebooks/v2/diagnose_instance.py new file mode 100644 index 00000000000..715f541c612 --- /dev/null +++ b/notebooks/v2/diagnose_instance.py @@ -0,0 +1,97 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_diagnose_instance] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def diagnose_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", + gcs_bucket: str = "gs://example-bucket", +) -> None: + """Runs diagnostics on a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to diagnose. + gcs_bucket: The Cloud Storage bucket URI to store diagnostic logs. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + diagnostic_config = notebooks_v2.DiagnosticConfig( + gcs_bucket=gcs_bucket, + relative_path="example-diagnostics", + enable_repair_flag=False, + ) + request = notebooks_v2.DiagnoseInstanceRequest( + name=name, + diagnostic_config=diagnostic_config, + ) + + try: + operation = client.diagnose_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"Health State: {response.health_state.name}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_diagnose_instance] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Run diagnostics on a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + parser.add_argument( + "--gcs-bucket", + default="gs://example-bucket", + help="The Cloud Storage bucket URI for diagnostic output.", + ) + args = parser.parse_args() + diagnose_instance( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + gcs_bucket=args.gcs_bucket, + ) diff --git a/notebooks/v2/get_instance.py b/notebooks/v2/get_instance.py new file mode 100644 index 00000000000..c66a4d417bc --- /dev/null +++ b/notebooks/v2/get_instance.py @@ -0,0 +1,81 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_get_instance] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def get_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Gets details for a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to retrieve. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.GetInstanceRequest(name=name) + + try: + response = client.get_instance(request=request) + print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") + print(f"Health State: {response.health_state.name}") + print(f"Proxy URI: {response.proxy_uri}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_get_instance] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Get details for a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + get_instance( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/get_instance_iam_policy.py b/notebooks/v2/get_instance_iam_policy.py new file mode 100644 index 00000000000..2515e9fdfe4 --- /dev/null +++ b/notebooks/v2/get_instance_iam_policy.py @@ -0,0 +1,82 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_get_instance_iam_policy] +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.iam.v1 import iam_policy_pb2 + + +def get_instance_iam_policy( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Gets the IAM policy for a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = iam_policy_pb2.GetIamPolicyRequest(resource=name) + + try: + response = client.get_iam_policy(request=request) + print(f"Resource ID: {name}") + print(f"Policy Etag: {response.etag}") + print(f"Policy Version: {response.version}") + for binding in response.bindings: + print(f"Role: {binding.role}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_get_instance_iam_policy] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Get the IAM policy for a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + get_instance_iam_policy( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/instances_test.py b/notebooks/v2/instances_test.py new file mode 100644 index 00000000000..1b52edd5acf --- /dev/null +++ b/notebooks/v2/instances_test.py @@ -0,0 +1,343 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.iam.v1 import iam_policy_pb2 +from google.iam.v1 import policy_pb2 +import pytest + +import check_instance_upgradability +import create_instance +import create_instance_cmek +import create_instance_confidential +import create_instance_custom_container +import create_instance_gpu +import create_instance_hardened +import create_instance_private_vpc +import create_instance_single_user +import delete_instance +import diagnose_instance +import get_instance +import get_instance_iam_policy +import list_instances +import reset_instance +import rollback_instance +import set_instance_iam_policy +import start_instance +import stop_instance +import test_instance_iam_permissions +import update_instance_machine_type +import update_instance_metadata +import update_instance_shielded_config +import upgrade_instance + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" +) + + +def _make_mock_instance( + name: str = INSTANCE_NAME, + machine_type: str = "e2-standard-4", +) -> notebooks_v2.Instance: + return notebooks_v2.Instance( + name=name, + id="1234567890", + creator="example-user@example.com", + proxy_uri="example-proxy.notebooks.googleusercontent.com", + state=notebooks_v2.State.ACTIVE, + health_state=notebooks_v2.HealthState.HEALTHY, + gce_setup=notebooks_v2.GceSetup( + machine_type=machine_type, + disable_public_ip=True, + metadata={"idle-timeout-seconds": "3600"}, + boot_disk=notebooks_v2.BootDisk( + disk_size_gb=200, + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=( + f"projects/{PROJECT_ID}/locations/us-central1/" + "keyRings/example-key-ring/cryptoKeys/example-key" + ), + ), + container_image=notebooks_v2.ContainerImage( + repository="gcr.io/deeplearning-platform-release/base-cpu", + tag="latest", + ), + shielded_instance_config=notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ), + ), + ) + + +def _make_mock_lro(result_obj: object = None) -> mock.MagicMock: + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = result_obj + return operation + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_list_instances_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.list_instances.return_value = [_make_mock_instance()] + + list_instances.list_instances(PROJECT_ID, LOCATION_ID) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_list_instances_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.list_instances.side_effect = exceptions.NotFound("Zone not found") + + list_instances.list_instances(PROJECT_ID, LOCATION_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert LOCATION_ID in captured.err + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_get_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_instance.return_value = _make_mock_instance() + + get_instance.get_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Health State: HEALTHY" in captured.out + assert "State: ACTIVE" in captured.out + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_get_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_instance.side_effect = exceptions.NotFound("Instance missing") + + get_instance.get_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.return_value = _make_mock_lro( + _make_mock_instance() + ) + + create_instance.create_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Machine Type: e2-standard-4" in captured.out + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Duplicate instance" + ) + + create_instance.create_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_stop_start_reset_delete_instance( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.stop_instance.return_value = _make_mock_lro( + _make_mock_instance() + ) + mock_client.start_instance.return_value = _make_mock_lro( + _make_mock_instance() + ) + mock_client.reset_instance.return_value = _make_mock_lro( + _make_mock_instance() + ) + mock_client.delete_instance.return_value = _make_mock_lro(None) + + stop_instance.stop_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + start_instance.start_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + reset_instance.reset_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + delete_instance.delete_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Status: Deleted" in captured.out + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_update_and_maintenance_operations( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.update_instance.return_value = _make_mock_lro( + _make_mock_instance(machine_type="n1-standard-8") + ) + mock_client.check_instance_upgradability.return_value = ( + notebooks_v2.CheckInstanceUpgradabilityResponse( + upgradeable=True, + upgrade_version="m120", + upgrade_info="Upgrade available", + upgrade_image="gcr.io/deeplearning-platform-release/base-cpu:m120", + ) + ) + mock_client.upgrade_instance.return_value = _make_mock_lro( + _make_mock_instance() + ) + mock_client.rollback_instance.return_value = _make_mock_lro( + _make_mock_instance() + ) + mock_client.diagnose_instance.return_value = _make_mock_lro( + _make_mock_instance() + ) + + update_instance_machine_type.update_instance_machine_type( + PROJECT_ID, LOCATION_ID, INSTANCE_ID, "n1-standard-8" + ) + update_instance_metadata.update_instance_metadata( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + update_instance_shielded_config.update_instance_shielded_config( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + check_instance_upgradability.check_instance_upgradability( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + upgrade_instance.upgrade_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + rollback_instance.rollback_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + diagnose_instance.diagnose_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "Machine Type: n1-standard-8" in captured.out + assert "Secure Boot Enabled: True" in captured.out + assert "Upgradeable: True" in captured.out + assert "Health State: HEALTHY" in captured.out + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_specialized_create_and_iam_operations( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.return_value = _make_mock_lro( + _make_mock_instance() + ) + mock_policy = policy_pb2.Policy( + version=1, + etag=b"etag-1", + bindings=[ + policy_pb2.Binding( + role="roles/notebooks.runner", + members=["user:example-user@example.com"], + ) + ], + ) + mock_client.get_iam_policy.return_value = mock_policy + mock_client.set_iam_policy.return_value = mock_policy + mock_client.test_iam_permissions.return_value = ( + iam_policy_pb2.TestIamPermissionsResponse( + permissions=["notebooks.instances.get"] + ) + ) + + create_instance_gpu.create_instance_gpu( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + create_instance_custom_container.create_instance_custom_container( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + create_instance_cmek.create_instance_cmek( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + create_instance_confidential.create_instance_confidential( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + create_instance_private_vpc.create_instance_private_vpc( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + create_instance_single_user.create_instance_single_user( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + create_instance_hardened.create_instance_hardened( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + get_instance_iam_policy.get_instance_iam_policy( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + set_instance_iam_policy.set_instance_iam_policy( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + test_instance_iam_permissions.test_instance_iam_permissions( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert ( + "Container Repository: gcr.io/deeplearning-platform-release/base-cpu" + in captured.out + ) + assert "Boot Disk Encryption: CMEK" in captured.out + assert "Disable Public IP: True" in captured.out + assert "Creator: example-user@example.com" in captured.out + assert "Secure Boot Enabled: True" in captured.out + assert "Role: roles/notebooks.runner" in captured.out + assert "Allowed Permissions: ['notebooks.instances.get']" in captured.out diff --git a/notebooks/v2/list_instances.py b/notebooks/v2/list_instances.py new file mode 100644 index 00000000000..84f11113160 --- /dev/null +++ b/notebooks/v2/list_instances.py @@ -0,0 +1,74 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_list_instances] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def list_instances( + project_id: str, + location_id: str = "us-central1-a", +) -> None: + """Lists Workbench instances in a zone. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone to list instances from. + """ + client = notebooks_v2.NotebookServiceClient() + parent = f"projects/{project_id}/locations/{location_id}" + request = notebooks_v2.ListInstancesRequest(parent=parent) + + try: + response = client.list_instances(request=request) + for instance in response: + print(f"Resource ID: {instance.name}") + print(f"Creation Date: {instance.create_time}") + print(f"Instance ID: {instance.id}") + print(f"Proxy URI: {instance.proxy_uri}") + print(f"State: {instance.state.name}") + except exceptions.NotFound as err: + print( + f"error: Could not list instances in location '{location_id}' " + f"for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_list_instances] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="List Workbench instances." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + args = parser.parse_args() + list_instances( + project_id=args.project_id, + location_id=args.location_id, + ) diff --git a/notebooks/v2/noxfile_config.py b/notebooks/v2/noxfile_config.py new file mode 100644 index 00000000000..719a62bc276 --- /dev/null +++ b/notebooks/v2/noxfile_config.py @@ -0,0 +1,31 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +# Default TEST_CONFIG_OVERRIDE for python repos. + +TEST_CONFIG_OVERRIDE = { + # You can opt out from the test for specific Python versions. + "ignored_versions": [], + # Old samples are opted out of enforcing Python type hints + # All new samples should feature them + "enforce_type_hints": True, + # An envvar key for determining the project id to use. + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + # If you need to use a specific version of pip, + # change pip_version_override to the string representation + # of the version number, for example, "20.2.4" + "pip_version_override": None, + # A dictionary you want to inject into your test. + "envs": {}, +} diff --git a/notebooks/v2/requirements-test.txt b/notebooks/v2/requirements-test.txt new file mode 100644 index 00000000000..b3b2b3ca4cc --- /dev/null +++ b/notebooks/v2/requirements-test.txt @@ -0,0 +1 @@ +pytest==9.0.3 diff --git a/notebooks/v2/requirements.txt b/notebooks/v2/requirements.txt new file mode 100644 index 00000000000..877976ff7b1 --- /dev/null +++ b/notebooks/v2/requirements.txt @@ -0,0 +1 @@ +google-cloud-notebooks==1.18.0 diff --git a/notebooks/v2/reset_instance.py b/notebooks/v2/reset_instance.py new file mode 100644 index 00000000000..2c95aa60be8 --- /dev/null +++ b/notebooks/v2/reset_instance.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_reset_instance] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def reset_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Resets a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to reset. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.ResetInstanceRequest(name=name) + + try: + operation = client.reset_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_reset_instance] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Reset a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + reset_instance( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/rollback_instance.py b/notebooks/v2/rollback_instance.py new file mode 100644 index 00000000000..696ef80dd9c --- /dev/null +++ b/notebooks/v2/rollback_instance.py @@ -0,0 +1,101 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_rollback_instance] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def rollback_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", + snapshot_id: str = "example-snapshot", + revision_id: str = "1", +) -> None: + """Rolls back a Workbench instance to a snapshot. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to roll back. + snapshot_id: The ID of the Compute Engine snapshot in the project. + revision_id: The revision ID to roll back to. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + target_snapshot = f"projects/{project_id}/global/snapshots/{snapshot_id}" + request = notebooks_v2.RollbackInstanceRequest( + name=name, + target_snapshot=target_snapshot, + revision_id=revision_id, + ) + + try: + operation = client.rollback_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_rollback_instance] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Roll back a Workbench instance to a snapshot." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + parser.add_argument( + "--snapshot-id", + default="example-snapshot", + help="The ID of the target Compute Engine snapshot.", + ) + parser.add_argument( + "--revision-id", + default="1", + help="The revision ID for rollback.", + ) + args = parser.parse_args() + rollback_instance( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + snapshot_id=args.snapshot_id, + revision_id=args.revision_id, + ) diff --git a/notebooks/v2/set_instance_iam_policy.py b/notebooks/v2/set_instance_iam_policy.py new file mode 100644 index 00000000000..f1471aceb5f --- /dev/null +++ b/notebooks/v2/set_instance_iam_policy.py @@ -0,0 +1,107 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_set_instance_iam_policy] +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.iam.v1 import iam_policy_pb2 +from google.iam.v1 import policy_pb2 + + +def set_instance_iam_policy( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", + member: str = "user:example-user@example.com", + role: str = "roles/notebooks.runner", +) -> None: + """Sets the IAM policy for a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance. + member: The principal to grant the role to. + role: The IAM role to bind on the instance. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + binding = policy_pb2.Binding( + role=role, + members=[member], + ) + policy = policy_pb2.Policy(bindings=[binding]) + request = iam_policy_pb2.SetIamPolicyRequest( + resource=name, + policy=policy, + ) + + try: + response = client.set_iam_policy(request=request) + print(f"Resource ID: {name}") + print(f"Policy Etag: {response.etag}") + print(f"Policy Version: {response.version}") + for item in response.bindings: + print(f"Role: {item.role}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_set_instance_iam_policy] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Set the IAM policy for a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + parser.add_argument( + "--member", + default="user:example-user@example.com", + help="The IAM principal to grant access to.", + ) + parser.add_argument( + "--role", + default="roles/notebooks.runner", + help="The IAM role to grant.", + ) + args = parser.parse_args() + set_instance_iam_policy( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + member=args.member, + role=args.role, + ) diff --git a/notebooks/v2/start_instance.py b/notebooks/v2/start_instance.py new file mode 100644 index 00000000000..84ae2c4f868 --- /dev/null +++ b/notebooks/v2/start_instance.py @@ -0,0 +1,81 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_start_instance] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def start_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Starts a stopped Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to start. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.StartInstanceRequest(name=name) + + try: + operation = client.start_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"Proxy URI: {response.proxy_uri}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_start_instance] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Start a stopped Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + start_instance( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/stop_instance.py b/notebooks/v2/stop_instance.py new file mode 100644 index 00000000000..2b0d744b091 --- /dev/null +++ b/notebooks/v2/stop_instance.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_stop_instance] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def stop_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Stops a running Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to stop. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.StopInstanceRequest(name=name) + + try: + operation = client.stop_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_stop_instance] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Stop a running Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + stop_instance( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/test_instance_iam_permissions.py b/notebooks/v2/test_instance_iam_permissions.py new file mode 100644 index 00000000000..53fb800daa4 --- /dev/null +++ b/notebooks/v2/test_instance_iam_permissions.py @@ -0,0 +1,82 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_test_instance_iam_permissions] +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.iam.v1 import iam_policy_pb2 + + +def test_instance_iam_permissions( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Tests IAM permissions on a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = iam_policy_pb2.TestIamPermissionsRequest( + resource=name, + permissions=["notebooks.instances.get"], + ) + + try: + response = client.test_iam_permissions(request=request) + print(f"Resource ID: {name}") + print(f"Allowed Permissions: {list(response.permissions)}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_test_instance_iam_permissions] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Test IAM permissions on a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + test_instance_iam_permissions( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/update_instance_machine_type.py b/notebooks/v2/update_instance_machine_type.py new file mode 100644 index 00000000000..59f34293551 --- /dev/null +++ b/notebooks/v2/update_instance_machine_type.py @@ -0,0 +1,96 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_update_instance_machine_type] +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.protobuf import field_mask_pb2 + + +def update_instance_machine_type( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", + machine_type: str = "n1-standard-8", +) -> None: + """Updates the machine type of a stopped Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the stopped Workbench instance to update. + machine_type: The target Compute Engine machine type. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + gce_setup = notebooks_v2.GceSetup(machine_type=machine_type) + instance = notebooks_v2.Instance(name=name, gce_setup=gce_setup) + update_mask = field_mask_pb2.FieldMask(paths=["gce_setup.machine_type"]) + request = notebooks_v2.UpdateInstanceRequest( + instance=instance, + update_mask=update_mask, + ) + + try: + operation = client.update_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"Machine Type: {response.gce_setup.machine_type}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_update_instance_machine_type] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Update the machine type of a stopped Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + parser.add_argument( + "--machine-type", + default="n1-standard-8", + help="The Compute Engine machine type.", + ) + args = parser.parse_args() + update_instance_machine_type( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + machine_type=args.machine_type, + ) diff --git a/notebooks/v2/update_instance_metadata.py b/notebooks/v2/update_instance_metadata.py new file mode 100644 index 00000000000..f0f6c19706b --- /dev/null +++ b/notebooks/v2/update_instance_metadata.py @@ -0,0 +1,90 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_update_instance_metadata] +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.protobuf import field_mask_pb2 + + +def update_instance_metadata( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Updates custom metadata on a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to update. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + gce_setup = notebooks_v2.GceSetup( + metadata={"idle-timeout-seconds": "3600"}, + ) + instance = notebooks_v2.Instance(name=name, gce_setup=gce_setup) + update_mask = field_mask_pb2.FieldMask(paths=["gce_setup.metadata"]) + request = notebooks_v2.UpdateInstanceRequest( + instance=instance, + update_mask=update_mask, + ) + + try: + operation = client.update_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"Metadata: {response.gce_setup.metadata}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_update_instance_metadata] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Update custom metadata on a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + update_instance_metadata( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/update_instance_shielded_config.py b/notebooks/v2/update_instance_shielded_config.py new file mode 100644 index 00000000000..f9a45c63306 --- /dev/null +++ b/notebooks/v2/update_instance_shielded_config.py @@ -0,0 +1,102 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_update_instance_shielded_config] +from google.api_core import exceptions +from google.cloud import notebooks_v2 +from google.protobuf import field_mask_pb2 + + +def update_instance_shielded_config( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Updates the Shielded VM configuration on a stopped Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the stopped Workbench instance to update. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + shielded_config = notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ) + gce_setup = notebooks_v2.GceSetup(shielded_instance_config=shielded_config) + instance = notebooks_v2.Instance(name=name, gce_setup=gce_setup) + update_mask = field_mask_pb2.FieldMask( + paths=[ + "gce_setup.shielded_instance_config.enable_secure_boot", + "gce_setup.shielded_instance_config.enable_vtpm", + "gce_setup.shielded_instance_config.enable_integrity_monitoring", + ] + ) + request = notebooks_v2.UpdateInstanceRequest( + instance=instance, + update_mask=update_mask, + ) + + try: + operation = client.update_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + secure_boot = ( + response.gce_setup.shielded_instance_config.enable_secure_boot + ) + print(f"Resource ID: {response.name}") + print(f"Secure Boot Enabled: {secure_boot}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_update_instance_shielded_config] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Update the Shielded VM config of a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + update_instance_shielded_config( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) diff --git a/notebooks/v2/upgrade_instance.py b/notebooks/v2/upgrade_instance.py new file mode 100644 index 00000000000..9e24791c40c --- /dev/null +++ b/notebooks/v2/upgrade_instance.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +import argparse +import sys + +# [START notebooks_v2_upgrade_instance] +from google.api_core import exceptions +from google.cloud import notebooks_v2 + + +def upgrade_instance( + project_id: str, + location_id: str = "us-central1-a", + instance_id: str = "example-instance", +) -> None: + """Upgrades a Workbench instance. + + Args: + project_id: The Google Cloud project ID. + location_id: The Compute Engine zone where the instance resides. + instance_id: The ID of the Workbench instance to upgrade. + """ + client = notebooks_v2.NotebookServiceClient() + name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + request = notebooks_v2.UpgradeInstanceRequest(name=name) + + try: + operation = client.upgrade_instance(request=request) + print(f"Waiting for operation to complete: {operation.operation.name}") + response = operation.result() + print(f"Resource ID: {response.name}") + print(f"State: {response.state.name}") + except exceptions.NotFound as err: + print( + f"error: Instance '{instance_id}' not found in location " + f"'{location_id}' for project '{project_id}': {err.message}", + file=sys.stderr, + ) + + +# [END notebooks_v2_upgrade_instance] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Upgrade a Workbench instance." + ) + parser.add_argument( + "--project-id", + required=True, + help="The Google Cloud project ID.", + ) + parser.add_argument( + "--location-id", + default="us-central1-a", + help="The Compute Engine zone.", + ) + parser.add_argument( + "--instance-id", + default="example-instance", + help="The ID of the Workbench instance.", + ) + args = parser.parse_args() + upgrade_instance( + project_id=args.project_id, + location_id=args.location_id, + instance_id=args.instance_id, + ) From 7ce80937f63a1b92a42d638965714b221960d477 Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Tue, 3 Nov 2026 18:14:44 +0000 Subject: [PATCH 2/8] refactor(notebooks): align v2 samples with authoring standards, add per-sample tests, and use IAM read-modify-write Problem: 1. The initial Notebooks API v2 samples included inline argparse / __main__ blocks inside the sample files, placed region tags after import sys, hardcoded example-project in resource paths, and used a single combined instances_test.py file. 2. test_instance_iam_permissions.py started with 'test_', causing pytest to collect the sample function itself as a test during CI runs. 3. set_instance_iam_policy.py constructed a new Policy object with a single binding instead of reading and modifying the existing IAM policy. 4. notebooks/v2/noxfile_config.py did not align ignored_versions with the repository root noxfile_config.py. Solution: 1. Refactor all 24 samples in notebooks/v2/ to pure functions enclosed by region tags at the top of each file, consistently referring to 'Workbench instance' in comments/docstrings and interpolating project_id across all resource URIs. 2. Rename test_instance_iam_permissions.py to check_instance_iam_permissions.py (notebooks_v2_check_instance_iam_permissions) and split instances_test.py into 24 separate _test.py files (48 unit tests total). 3. Update set_instance_iam_policy.py to use the read-modify-write pattern (get_iam_policy -> policy.bindings.append -> set_iam_policy) and harden region_id extraction in create_instance_cmek.py, create_instance_private_vpc.py, and create_instance_hardened.py. 4. Update notebooks/README.md and notebooks/v2/README.md with enterprise security and administration positioning for Workbench instances, and align notebooks/v2/noxfile_config.py with the root configuration. --- notebooks/README.md | 74 +--- notebooks/v2/README.md | 12 +- ...s.py => check_instance_iam_permissions.py} | 36 +- .../v2/check_instance_iam_permissions_test.py | 67 ++++ notebooks/v2/check_instance_upgradability.py | 30 +- .../v2/check_instance_upgradability_test.py | 70 ++++ notebooks/v2/create_instance.py | 30 +- notebooks/v2/create_instance_cmek.py | 57 +-- notebooks/v2/create_instance_cmek_test.py | 77 ++++ notebooks/v2/create_instance_confidential.py | 34 +- .../v2/create_instance_confidential_test.py | 80 ++++ .../v2/create_instance_custom_container.py | 41 +-- .../create_instance_custom_container_test.py | 81 +++++ notebooks/v2/create_instance_gpu.py | 30 +- notebooks/v2/create_instance_gpu_test.py | 69 ++++ notebooks/v2/create_instance_hardened.py | 94 +++-- notebooks/v2/create_instance_hardened_test.py | 84 +++++ notebooks/v2/create_instance_private_vpc.py | 59 +-- .../v2/create_instance_private_vpc_test.py | 76 ++++ notebooks/v2/create_instance_single_user.py | 46 +-- .../v2/create_instance_single_user_test.py | 74 ++++ notebooks/v2/create_instance_test.py | 71 ++++ notebooks/v2/delete_instance.py | 32 +- notebooks/v2/delete_instance_test.py | 63 ++++ notebooks/v2/diagnose_instance.py | 36 +- notebooks/v2/diagnose_instance_test.py | 69 ++++ notebooks/v2/get_instance.py | 30 +- notebooks/v2/get_instance_iam_policy.py | 30 +- notebooks/v2/get_instance_iam_policy_test.py | 70 ++++ notebooks/v2/get_instance_test.py | 66 ++++ notebooks/v2/instances_test.py | 343 ------------------ notebooks/v2/list_instances.py | 24 +- notebooks/v2/list_instances_test.py | 65 ++++ notebooks/v2/noxfile_config.py | 4 +- notebooks/v2/reset_instance.py | 30 +- notebooks/v2/reset_instance_test.py | 66 ++++ notebooks/v2/rollback_instance.py | 42 +-- notebooks/v2/rollback_instance_test.py | 67 ++++ notebooks/v2/set_instance_iam_policy.py | 65 +--- notebooks/v2/set_instance_iam_policy_test.py | 79 ++++ notebooks/v2/start_instance.py | 30 +- notebooks/v2/start_instance_test.py | 67 ++++ notebooks/v2/stop_instance.py | 30 +- notebooks/v2/stop_instance_test.py | 66 ++++ notebooks/v2/update_instance_machine_type.py | 36 +- .../v2/update_instance_machine_type_test.py | 71 ++++ notebooks/v2/update_instance_metadata.py | 30 +- notebooks/v2/update_instance_metadata_test.py | 71 ++++ .../v2/update_instance_shielded_config.py | 34 +- .../update_instance_shielded_config_test.py | 77 ++++ notebooks/v2/upgrade_instance.py | 30 +- notebooks/v2/upgrade_instance_test.py | 65 ++++ 52 files changed, 1854 insertions(+), 1226 deletions(-) rename notebooks/v2/{test_instance_iam_permissions.py => check_instance_iam_permissions.py} (64%) create mode 100644 notebooks/v2/check_instance_iam_permissions_test.py create mode 100644 notebooks/v2/check_instance_upgradability_test.py create mode 100644 notebooks/v2/create_instance_cmek_test.py create mode 100644 notebooks/v2/create_instance_confidential_test.py create mode 100644 notebooks/v2/create_instance_custom_container_test.py create mode 100644 notebooks/v2/create_instance_gpu_test.py create mode 100644 notebooks/v2/create_instance_hardened_test.py create mode 100644 notebooks/v2/create_instance_private_vpc_test.py create mode 100644 notebooks/v2/create_instance_single_user_test.py create mode 100644 notebooks/v2/create_instance_test.py create mode 100644 notebooks/v2/delete_instance_test.py create mode 100644 notebooks/v2/diagnose_instance_test.py create mode 100644 notebooks/v2/get_instance_iam_policy_test.py create mode 100644 notebooks/v2/get_instance_test.py delete mode 100644 notebooks/v2/instances_test.py create mode 100644 notebooks/v2/list_instances_test.py create mode 100644 notebooks/v2/reset_instance_test.py create mode 100644 notebooks/v2/rollback_instance_test.py create mode 100644 notebooks/v2/set_instance_iam_policy_test.py create mode 100644 notebooks/v2/start_instance_test.py create mode 100644 notebooks/v2/stop_instance_test.py create mode 100644 notebooks/v2/update_instance_machine_type_test.py create mode 100644 notebooks/v2/update_instance_metadata_test.py create mode 100644 notebooks/v2/update_instance_shielded_config_test.py create mode 100644 notebooks/v2/upgrade_instance_test.py diff --git a/notebooks/README.md b/notebooks/README.md index 457942edab4..845a421ba58 100644 --- a/notebooks/README.md +++ b/notebooks/README.md @@ -1,70 +1,12 @@ -# Workbench Instances and Notebooks +# Google Cloud Notebooks & Workbench Samples -**Workbench instances** provide a managed, integrated, and secure JupyterLab environment for data scientists, machine learning developers, and agent builders to experiment, develop, and deploy workloads on Google Cloud. +This directory contains samples and tutorials for managing **Workbench instances** and executing Jupyter notebooks on Google Cloud. -This directory contains Python client library code samples and Jupyter notebook tutorials: -- [Notebooks API v2 code samples](v2) -- [Jupyter notebook tutorials](tutorials) +Workbench instances provide enterprises with a managed, secure JupyterLab environment that protects organizational data while enabling centralized administration and configuration. Data scientists can experiment, develop, and deploy machine learning and analytics workloads on Google Cloud while maintaining their organization's required security posture. -## Notebooks API v2 (`notebooks/v2`) +## Directory Structure -In the [v2](v2) folder, you will find atomic Python code samples using the `google-cloud-notebooks` (`google.cloud.notebooks_v2`) client library to create, manage, configure, upgrade, and diagnose Workbench instances. - -## Tutorials (`notebooks/tutorials`) - -### Run in a Workbench Instance - -1. Create a Workbench instance using the Google Cloud console, the `gcloud workbench instances create` CLI command, or the [create_instance.py](v2/create_instance.py) sample. -2. Clone this repository via the JupyterLab terminal or using the Git clone button: - -```bash -git clone https://github.com/GoogleCloudPlatform/python-docs-samples.git -``` - -3. Open the `notebooks/tutorials` folder and launch the Jupyter notebooks. - -### Run Locally - -1. Install JupyterLab or Jupyter Notebook ([installation instructions](https://jupyter.org/install)). -2. Install the dependencies in the [requirements.txt](./requirements.txt) file ([instructions below](#install-the-dependencies)). -3. Register the `google-cloud-bigquery` magic commands ([instructions below](#register-magics-and-configure-matplotlib)). -4. Set `matplotlib` to render inline ([instructions below](#register-magics-and-configure-matplotlib)). - -#### Install the dependencies - -Install the dependencies with the following command: - -```bash -pip install --upgrade -r requirements.txt -``` - -#### Register magics and configure matplotlib - -You can either perform these setup steps in a single notebook, or add the steps to your IPython configuration file to apply to all notebooks. - -##### Perform setup steps within a notebook - -Run the following commands in your notebook to register the BigQuery magic commands and set `matplotlib` to render inline: - -```python -%load_ext google.cloud.bigquery -%matplotlib inline -``` - -##### Perform setup steps in your IPython configuration file - -Add the following code to your `ipython_config.py` file to register the BigQuery magic commands and set `matplotlib` to render inline across all notebooks: - -```python -c = get_config() - -# Register magic commands -c.InteractiveShellApp.extensions = [ - "google.cloud.bigquery", -] - -# Enable matplotlib renderings to render inline in the notebook. -c.InteractiveShellApp.matplotlib = "inline" -``` - -See the [IPython configuration documentation](https://ipython.readthedocs.io/en/stable/config/intro.html) for more information. +- **[`v2/`](v2/)**: Atomic Python client library samples (`google-cloud-notebooks` / `google.cloud.notebooks_v2`) for creating, configuring, hardening, and managing **Workbench instances**, including enterprise security controls (private VPC networking, single-user least-privilege service accounts, CMEK encryption, Shielded VM configuration, and IAM access policies). +- **[`rendered/`](rendered/)**: Rendered HTML views of tutorial notebooks. +- **[`tutorials/bigquery/`](tutorials/bigquery/)**: Jupyter notebook tutorials demonstrating BigQuery integration from a notebook environment. +- **[`tutorials/vision/`](tutorials/vision/)**: Jupyter notebook tutorials demonstrating Cloud Vision API integration from a notebook environment. diff --git a/notebooks/v2/README.md b/notebooks/v2/README.md index fb44cca16ba..7b8631fb033 100644 --- a/notebooks/v2/README.md +++ b/notebooks/v2/README.md @@ -2,6 +2,8 @@ This directory contains atomic Python code samples demonstrating how to manage **Workbench instances** using the `google-cloud-notebooks` (`google.cloud.notebooks_v2`) client library. +Workbench instances target enterprises that want to manage a secure JupyterLab environment that keeps their data safe while allowing centralized administration and configuration. Data scientists can experiment, develop, and deploy workloads on Google Cloud while maintaining their organization's desired security posture. + ## Prerequisites 1. Enable the **Notebooks API** (`notebooks.googleapis.com`) in your Google Cloud project: @@ -29,7 +31,7 @@ This directory contains atomic Python code samples demonstrating how to manage * - [`reset_instance.py`](reset_instance.py) — Resets (reboots) a Workbench instance (`notebooks_v2_reset_instance`). - [`delete_instance.py`](delete_instance.py) — Deletes a Workbench instance (`notebooks_v2_delete_instance`). -### Instance Configuration, Maintenance & Diagnostics +### Centralized Administration, Maintenance & Diagnostics - [`update_instance_machine_type.py`](update_instance_machine_type.py) — Updates the Compute Engine machine type of a stopped Workbench instance (`notebooks_v2_update_instance_machine_type`). - [`update_instance_metadata.py`](update_instance_metadata.py) — Updates custom metadata keys on a Workbench instance (`notebooks_v2_update_instance_metadata`). - [`update_instance_shielded_config.py`](update_instance_shielded_config.py) — Updates the Shielded VM configuration on a stopped Workbench instance (`notebooks_v2_update_instance_shielded_config`). @@ -38,7 +40,7 @@ This directory contains atomic Python code samples demonstrating how to manage * - [`rollback_instance.py`](rollback_instance.py) — Rolls back a Workbench instance to a previous snapshot and revision (`notebooks_v2_rollback_instance`). - [`diagnose_instance.py`](diagnose_instance.py) — Runs a diagnostic bundle export to Cloud Storage for a Workbench instance (`notebooks_v2_diagnose_instance`). -### Specialized Creation, Security Blueprint Hardening & Access Control +### Enterprise Security Posture, Specialized Compute & Access Control - [`create_instance_gpu.py`](create_instance_gpu.py) — Creates a GPU-accelerated Workbench instance (`notebooks_v2_create_instance_gpu`). - [`create_instance_custom_container.py`](create_instance_custom_container.py) — Creates a Workbench instance from a custom container image (`notebooks_v2_create_instance_custom_container`). - [`create_instance_cmek.py`](create_instance_cmek.py) — Creates a Workbench instance encrypted with a customer-managed encryption key (`notebooks_v2_create_instance_cmek`). @@ -47,12 +49,12 @@ This directory contains atomic Python code samples demonstrating how to manage * - [`create_instance_single_user.py`](create_instance_single_user.py) — Creates a single-user Workbench instance with a custom least-privilege service account (`notebooks_v2_create_instance_single_user`). - [`create_instance_hardened.py`](create_instance_hardened.py) — Creates a security-hardened Workbench instance disabling root, downloads, terminal, and public IP (`notebooks_v2_create_instance_hardened`). - [`get_instance_iam_policy.py`](get_instance_iam_policy.py) — Gets the IAM policy for a Workbench instance (`notebooks_v2_get_instance_iam_policy`). -- [`set_instance_iam_policy.py`](set_instance_iam_policy.py) — Sets an IAM policy binding on a Workbench instance (`notebooks_v2_set_instance_iam_policy`). -- [`test_instance_iam_permissions.py`](test_instance_iam_permissions.py) — Tests IAM permissions on a Workbench instance (`notebooks_v2_test_instance_iam_permissions`). +- [`set_instance_iam_policy.py`](set_instance_iam_policy.py) — Sets an IAM policy binding on a Workbench instance using read-modify-write (`notebooks_v2_set_instance_iam_policy`). +- [`check_instance_iam_permissions.py`](check_instance_iam_permissions.py) — Checks IAM permissions on a Workbench instance (`notebooks_v2_check_instance_iam_permissions`). ## Running Tests ```bash .venv/bin/pip install -r requirements-test.txt -.venv/bin/pytest instances_test.py +.venv/bin/pytest ``` diff --git a/notebooks/v2/test_instance_iam_permissions.py b/notebooks/v2/check_instance_iam_permissions.py similarity index 64% rename from notebooks/v2/test_instance_iam_permissions.py rename to notebooks/v2/check_instance_iam_permissions.py index 53fb800daa4..d43de2d2253 100644 --- a/notebooks/v2/test_instance_iam_permissions.py +++ b/notebooks/v2/check_instance_iam_permissions.py @@ -12,21 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_check_instance_iam_permissions] import sys -# [START notebooks_v2_test_instance_iam_permissions] from google.api_core import exceptions from google.cloud import notebooks_v2 from google.iam.v1 import iam_policy_pb2 -def test_instance_iam_permissions( +def check_instance_iam_permissions( project_id: str, location_id: str = "us-central1-a", instance_id: str = "example-instance", ) -> None: - """Tests IAM permissions on a Workbench instance. + """Checks IAM permissions on a Workbench instance. Args: project_id: The Google Cloud project ID. @@ -52,31 +51,4 @@ def test_instance_iam_permissions( ) -# [END notebooks_v2_test_instance_iam_permissions] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Test IAM permissions on a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - test_instance_iam_permissions( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) +# [END notebooks_v2_check_instance_iam_permissions] diff --git a/notebooks/v2/check_instance_iam_permissions_test.py b/notebooks/v2/check_instance_iam_permissions_test.py new file mode 100644 index 00000000000..d3b0af85fa7 --- /dev/null +++ b/notebooks/v2/check_instance_iam_permissions_test.py @@ -0,0 +1,67 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.iam.v1 import iam_policy_pb2 +import pytest + +import check_instance_iam_permissions + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_check_instance_iam_permissions_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.test_iam_permissions.return_value = ( + iam_policy_pb2.TestIamPermissionsResponse( + permissions=["notebooks.instances.get"] + ) + ) + + check_instance_iam_permissions.check_instance_iam_permissions( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Allowed Permissions: ['notebooks.instances.get']" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_check_instance_iam_permissions_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.test_iam_permissions.side_effect = exceptions.NotFound( + "Instance not found" + ) + + check_instance_iam_permissions.check_instance_iam_permissions( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/check_instance_upgradability.py b/notebooks/v2/check_instance_upgradability.py index b41b19aac5d..c544c8c0df1 100644 --- a/notebooks/v2/check_instance_upgradability.py +++ b/notebooks/v2/check_instance_upgradability.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_check_instance_upgradability] import sys -# [START notebooks_v2_check_instance_upgradability] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -54,30 +53,3 @@ def check_instance_upgradability( # [END notebooks_v2_check_instance_upgradability] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Check whether a Workbench instance can be upgraded." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - check_instance_upgradability( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/check_instance_upgradability_test.py b/notebooks/v2/check_instance_upgradability_test.py new file mode 100644 index 00000000000..672336c868c --- /dev/null +++ b/notebooks/v2/check_instance_upgradability_test.py @@ -0,0 +1,70 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import check_instance_upgradability + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_check_instance_upgradability_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.check_instance_upgradability.return_value = ( + notebooks_v2.CheckInstanceUpgradabilityResponse( + upgradeable=True, + upgrade_version="m120", + upgrade_info="Upgrade available", + upgrade_image="gcr.io/deeplearning-platform-release/base-cpu:m120", + ) + ) + + check_instance_upgradability.check_instance_upgradability( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Upgradeable: True" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_check_instance_upgradability_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.check_instance_upgradability.side_effect = exceptions.NotFound( + "Instance not found" + ) + + check_instance_upgradability.check_instance_upgradability( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/create_instance.py b/notebooks/v2/create_instance.py index 0218e742970..100edd83fb9 100644 --- a/notebooks/v2/create_instance.py +++ b/notebooks/v2/create_instance.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_create_instance] import sys -# [START notebooks_v2_create_instance] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -60,30 +59,3 @@ def create_instance( # [END notebooks_v2_create_instance] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Create a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID for the new Workbench instance.", - ) - args = parser.parse_args() - create_instance( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/create_instance_cmek.py b/notebooks/v2/create_instance_cmek.py index 033c7bf04fe..3891cbf07bb 100644 --- a/notebooks/v2/create_instance_cmek.py +++ b/notebooks/v2/create_instance_cmek.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_create_instance_cmek] import sys -# [START notebooks_v2_create_instance_cmek] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -27,29 +26,32 @@ def create_instance_cmek( key_ring_id: str = "example-key-ring", crypto_key_id: str = "example-key", ) -> None: - """Creates a CMEK-encrypted Workbench instance. + """Creates a Workbench instance encrypted with a customer-managed key. Args: project_id: The Google Cloud project ID. location_id: The Compute Engine zone for the new instance. instance_id: The ID to assign to the new Workbench instance. - key_ring_id: The Cloud KMS key ring ID in the same region. + key_ring_id: The Cloud KMS key ring ID in the instance region. crypto_key_id: The Cloud KMS crypto key ID used for disk encryption. """ client = notebooks_v2.NotebookServiceClient() parent = f"projects/{project_id}/locations/{location_id}" - region_id = "-".join(location_id.split("-")[:-1]) + parts = location_id.split("-") + region_id = "-".join(parts[:-1]) if len(parts) > 2 else location_id kms_key = ( f"projects/{project_id}/locations/{region_id}/" f"keyRings/{key_ring_id}/cryptoKeys/{crypto_key_id}" ) boot_disk = notebooks_v2.BootDisk( disk_size_gb=150, + disk_type=notebooks_v2.DiskType.PD_SSD, disk_encryption=notebooks_v2.DiskEncryption.CMEK, kms_key=kms_key, ) data_disk = notebooks_v2.DataDisk( disk_size_gb=100, + disk_type=notebooks_v2.DiskType.PD_SSD, disk_encryption=notebooks_v2.DiskEncryption.CMEK, kms_key=kms_key, ) @@ -69,10 +71,10 @@ def create_instance_cmek( operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") response = operation.result() - encryption = response.gce_setup.boot_disk.disk_encryption.name + boot_encryption = response.gce_setup.boot_disk.disk_encryption.name print(f"Resource ID: {response.name}") - print(f"Boot Disk Encryption: {encryption}") - print(f"KMS Key: {response.gce_setup.boot_disk.kms_key}") + print(f"Creation Date: {response.create_time}") + print(f"Boot Disk Encryption: {boot_encryption}") print(f"State: {response.state.name}") except exceptions.AlreadyExists as err: print( @@ -83,42 +85,3 @@ def create_instance_cmek( # [END notebooks_v2_create_instance_cmek] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Create a CMEK-encrypted Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-cmek-instance", - help="The ID for the new Workbench instance.", - ) - parser.add_argument( - "--key-ring-id", - default="example-key-ring", - help="The Cloud KMS key ring ID.", - ) - parser.add_argument( - "--crypto-key-id", - default="example-key", - help="The Cloud KMS crypto key ID.", - ) - args = parser.parse_args() - create_instance_cmek( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - key_ring_id=args.key_ring_id, - crypto_key_id=args.crypto_key_id, - ) diff --git a/notebooks/v2/create_instance_cmek_test.py b/notebooks/v2/create_instance_cmek_test.py new file mode 100644 index 00000000000..cfc4e28d874 --- /dev/null +++ b/notebooks/v2/create_instance_cmek_test.py @@ -0,0 +1,77 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import create_instance_cmek + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-cmek-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_cmek_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + boot_disk=notebooks_v2.BootDisk( + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=( + f"projects/{PROJECT_ID}/locations/us-central1/" + "keyRings/example-key-ring/cryptoKeys/example-key" + ), + ) + ), + ) + mock_client.create_instance.return_value = operation + + create_instance_cmek.create_instance_cmek(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Boot Disk Encryption: CMEK" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_cmek_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_cmek.create_instance_cmek(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/create_instance_confidential.py b/notebooks/v2/create_instance_confidential.py index 01eaa6aedfa..d88d79f3286 100644 --- a/notebooks/v2/create_instance_confidential.py +++ b/notebooks/v2/create_instance_confidential.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_create_instance_confidential] import sys -# [START notebooks_v2_create_instance_confidential] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -54,9 +53,7 @@ def create_instance_confidential( operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") response = operation.result() - secure_boot = ( - response.gce_setup.shielded_instance_config.enable_secure_boot - ) + secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot print(f"Resource ID: {response.name}") print(f"Machine Type: {response.gce_setup.machine_type}") print(f"Secure Boot Enabled: {secure_boot}") @@ -70,30 +67,3 @@ def create_instance_confidential( # [END notebooks_v2_create_instance_confidential] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Create a Shielded and Confidential VM Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-confidential-instance", - help="The ID for the new Workbench instance.", - ) - args = parser.parse_args() - create_instance_confidential( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/create_instance_confidential_test.py b/notebooks/v2/create_instance_confidential_test.py new file mode 100644 index 00000000000..f48d0092931 --- /dev/null +++ b/notebooks/v2/create_instance_confidential_test.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import create_instance_confidential + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-confidential-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_confidential_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + machine_type="n2d-standard-4", + shielded_instance_config=notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ), + ), + ) + mock_client.create_instance.return_value = operation + + create_instance_confidential.create_instance_confidential( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Secure Boot Enabled: True" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_confidential_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_confidential.create_instance_confidential( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/create_instance_custom_container.py b/notebooks/v2/create_instance_custom_container.py index 45ef91e9383..bec6cf469d7 100644 --- a/notebooks/v2/create_instance_custom_container.py +++ b/notebooks/v2/create_instance_custom_container.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_create_instance_custom_container] import sys -# [START notebooks_v2_create_instance_custom_container] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -55,10 +54,9 @@ def create_instance_custom_container( operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") response = operation.result() + repo = response.gce_setup.container_image.repository print(f"Resource ID: {response.name}") - print( - f"Container Repository: {response.gce_setup.container_image.repository}" - ) + print(f"Container Repository: {repo}") print(f"State: {response.state.name}") except exceptions.AlreadyExists as err: print( @@ -69,36 +67,3 @@ def create_instance_custom_container( # [END notebooks_v2_create_instance_custom_container] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Create a Workbench instance using a custom container image." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-container-instance", - help="The ID for the new Workbench instance.", - ) - parser.add_argument( - "--repository", - default="gcr.io/deeplearning-platform-release/base-cpu", - help="The container image repository URI.", - ) - args = parser.parse_args() - create_instance_custom_container( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - repository=args.repository, - ) diff --git a/notebooks/v2/create_instance_custom_container_test.py b/notebooks/v2/create_instance_custom_container_test.py new file mode 100644 index 00000000000..cda0484c0bb --- /dev/null +++ b/notebooks/v2/create_instance_custom_container_test.py @@ -0,0 +1,81 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import create_instance_custom_container + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-container-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_custom_container_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + container_image=notebooks_v2.ContainerImage( + repository="gcr.io/deeplearning-platform-release/base-cpu", + tag="latest", + ) + ), + ) + mock_client.create_instance.return_value = operation + + create_instance_custom_container.create_instance_custom_container( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert ( + "Container Repository: gcr.io/deeplearning-platform-release/base-cpu" + in captured.out + ) + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_custom_container_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_custom_container.create_instance_custom_container( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/create_instance_gpu.py b/notebooks/v2/create_instance_gpu.py index 9ee3e851274..76aad0cd1f7 100644 --- a/notebooks/v2/create_instance_gpu.py +++ b/notebooks/v2/create_instance_gpu.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_create_instance_gpu] import sys -# [START notebooks_v2_create_instance_gpu] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -70,30 +69,3 @@ def create_instance_gpu( # [END notebooks_v2_create_instance_gpu] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Create a GPU-accelerated Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-gpu-instance", - help="The ID for the new Workbench instance.", - ) - args = parser.parse_args() - create_instance_gpu( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/create_instance_gpu_test.py b/notebooks/v2/create_instance_gpu_test.py new file mode 100644 index 00000000000..1cc5d9de4aa --- /dev/null +++ b/notebooks/v2/create_instance_gpu_test.py @@ -0,0 +1,69 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import create_instance_gpu + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-gpu-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_gpu_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup(machine_type="n1-standard-4"), + ) + mock_client.create_instance.return_value = operation + + create_instance_gpu.create_instance_gpu(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Machine Type: n1-standard-4" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_gpu_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_gpu.create_instance_gpu(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/create_instance_hardened.py b/notebooks/v2/create_instance_hardened.py index 8c185c0364f..dc355c79214 100644 --- a/notebooks/v2/create_instance_hardened.py +++ b/notebooks/v2/create_instance_hardened.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_create_instance_hardened] import sys -# [START notebooks_v2_create_instance_hardened] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -24,34 +23,89 @@ def create_instance_hardened( project_id: str, location_id: str = "us-central1-a", instance_id: str = "example-hardened-instance", + network_id: str = "example-network", + subnet_id: str = "example-subnet", + service_account_id: str = "example-workbench-sa", + instance_owner: str = "example-user@example.com", + key_ring_id: str = "example-key-ring", + crypto_key_id: str = "example-key", ) -> None: """Creates a security-hardened Workbench instance. + Combines enterprise hardening controls for regulated environments: + disables public IP, restricts proxy access to a single user with a custom + service account, enables Shielded VM integrity checks, encrypts disks with + CMEK, and disables root access, file downloads, and the JupyterLab terminal + via instance metadata. + Args: project_id: The Google Cloud project ID. location_id: The Compute Engine zone for the new instance. instance_id: The ID to assign to the new Workbench instance. + network_id: The ID of the VPC network in the project. + subnet_id: The ID of the subnetwork in the instance region. + service_account_id: The ID of the custom service account in the project. + instance_owner: The single user email permitted to access the instance. + key_ring_id: The Cloud KMS key ring ID in the instance region. + crypto_key_id: The Cloud KMS crypto key ID used for disk encryption. """ client = notebooks_v2.NotebookServiceClient() parent = f"projects/{project_id}/locations/{location_id}" + parts = location_id.split("-") + region_id = "-".join(parts[:-1]) if len(parts) > 2 else location_id + network = f"projects/{project_id}/global/networks/{network_id}" + subnet = f"projects/{project_id}/regions/{region_id}/subnetworks/{subnet_id}" + service_account_email = f"{service_account_id}@{project_id}.iam.gserviceaccount.com" + kms_key = ( + f"projects/{project_id}/locations/{region_id}/" + f"keyRings/{key_ring_id}/cryptoKeys/{crypto_key_id}" + ) + network_interface = notebooks_v2.NetworkInterface( + network=network, + subnet=subnet, + nic_type=notebooks_v2.NetworkInterface.NicType.GVNIC, + ) + service_account = notebooks_v2.ServiceAccount( + email=service_account_email, + ) shielded_config = notebooks_v2.ShieldedInstanceConfig( enable_secure_boot=True, enable_vtpm=True, enable_integrity_monitoring=True, ) + boot_disk = notebooks_v2.BootDisk( + disk_size_gb=150, + disk_type=notebooks_v2.DiskType.PD_SSD, + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=kms_key, + ) + data_disk = notebooks_v2.DataDisk( + disk_size_gb=100, + disk_type=notebooks_v2.DiskType.PD_SSD, + disk_encryption=notebooks_v2.DiskEncryption.CMEK, + kms_key=kms_key, + ) hardened_metadata = { "notebook-disable-root": "true", "notebook-disable-downloads": "true", "notebook-disable-terminal": "true", - "serial-port-logging-enable": "false", + "serial-port-enable": "false", + "block-project-ssh-keys": "true", } gce_setup = notebooks_v2.GceSetup( machine_type="e2-standard-4", disable_public_ip=True, + network_interfaces=[network_interface], + service_accounts=[service_account], shielded_instance_config=shielded_config, + boot_disk=boot_disk, + data_disks=[data_disk], metadata=hardened_metadata, ) - instance = notebooks_v2.Instance(gce_setup=gce_setup) + instance = notebooks_v2.Instance( + gce_setup=gce_setup, + instance_owners=[instance_owner], + ) request = notebooks_v2.CreateInstanceRequest( parent=parent, instance_id=instance_id, @@ -62,10 +116,9 @@ def create_instance_hardened( operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") response = operation.result() - secure_boot = ( - response.gce_setup.shielded_instance_config.enable_secure_boot - ) + secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") print(f"Disable Public IP: {response.gce_setup.disable_public_ip}") print(f"Secure Boot Enabled: {secure_boot}") print(f"State: {response.state.name}") @@ -78,30 +131,3 @@ def create_instance_hardened( # [END notebooks_v2_create_instance_hardened] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Create a security-hardened Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-hardened-instance", - help="The ID for the new Workbench instance.", - ) - args = parser.parse_args() - create_instance_hardened( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/create_instance_hardened_test.py b/notebooks/v2/create_instance_hardened_test.py new file mode 100644 index 00000000000..8e338a053c4 --- /dev/null +++ b/notebooks/v2/create_instance_hardened_test.py @@ -0,0 +1,84 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import create_instance_hardened + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-hardened-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_hardened_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + disable_public_ip=True, + boot_disk=notebooks_v2.BootDisk( + disk_encryption=notebooks_v2.DiskEncryption.CMEK + ), + shielded_instance_config=notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ), + ), + ) + mock_client.create_instance.return_value = operation + + create_instance_hardened.create_instance_hardened( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Disable Public IP: True" in captured.out + assert "Secure Boot Enabled: True" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_hardened_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_hardened.create_instance_hardened( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/create_instance_private_vpc.py b/notebooks/v2/create_instance_private_vpc.py index 1130196f2e0..c8f7c8604b4 100644 --- a/notebooks/v2/create_instance_private_vpc.py +++ b/notebooks/v2/create_instance_private_vpc.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_create_instance_private_vpc] import sys -# [START notebooks_v2_create_instance_private_vpc] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -27,22 +26,25 @@ def create_instance_private_vpc( network_id: str = "example-network", subnet_id: str = "example-subnet", ) -> None: - """Creates a private-IP Workbench instance in a custom VPC. + """Creates a Workbench instance with private IP in a custom VPC subnet. + + Disabling the external public IP address isolates the Workbench instance + from the public internet while allowing private connectivity through a + custom VPC network and Cloud NAT or Private Google Access. Args: project_id: The Google Cloud project ID. location_id: The Compute Engine zone for the new instance. instance_id: The ID to assign to the new Workbench instance. network_id: The ID of the VPC network in the project. - subnet_id: The ID of the subnetwork in the same region. + subnet_id: The ID of the subnetwork in the instance region. """ client = notebooks_v2.NotebookServiceClient() parent = f"projects/{project_id}/locations/{location_id}" - region_id = "-".join(location_id.split("-")[:-1]) + parts = location_id.split("-") + region_id = "-".join(parts[:-1]) if len(parts) > 2 else location_id network = f"projects/{project_id}/global/networks/{network_id}" - subnet = ( - f"projects/{project_id}/regions/{region_id}/subnetworks/{subnet_id}" - ) + subnet = f"projects/{project_id}/regions/{region_id}/subnetworks/{subnet_id}" network_interface = notebooks_v2.NetworkInterface( network=network, subnet=subnet, @@ -65,7 +67,9 @@ def create_instance_private_vpc( print(f"Waiting for operation to complete: {operation.operation.name}") response = operation.result() print(f"Resource ID: {response.name}") + print(f"Creation Date: {response.create_time}") print(f"Disable Public IP: {response.gce_setup.disable_public_ip}") + print(f"Machine Type: {response.gce_setup.machine_type}") print(f"State: {response.state.name}") except exceptions.AlreadyExists as err: print( @@ -76,42 +80,3 @@ def create_instance_private_vpc( # [END notebooks_v2_create_instance_private_vpc] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Create a private-IP Workbench instance in a custom VPC." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-private-instance", - help="The ID for the new Workbench instance.", - ) - parser.add_argument( - "--network-id", - default="example-network", - help="The VPC network ID.", - ) - parser.add_argument( - "--subnet-id", - default="example-subnet", - help="The subnetwork ID.", - ) - args = parser.parse_args() - create_instance_private_vpc( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - network_id=args.network_id, - subnet_id=args.subnet_id, - ) diff --git a/notebooks/v2/create_instance_private_vpc_test.py b/notebooks/v2/create_instance_private_vpc_test.py new file mode 100644 index 00000000000..ecdaae661df --- /dev/null +++ b/notebooks/v2/create_instance_private_vpc_test.py @@ -0,0 +1,76 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import create_instance_private_vpc + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-private-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_private_vpc_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup( + machine_type="e2-standard-4", + disable_public_ip=True, + ), + ) + mock_client.create_instance.return_value = operation + + create_instance_private_vpc.create_instance_private_vpc( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Disable Public IP: True" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_private_vpc_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_private_vpc.create_instance_private_vpc( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/create_instance_single_user.py b/notebooks/v2/create_instance_single_user.py index eb5490d67e1..f5077073525 100644 --- a/notebooks/v2/create_instance_single_user.py +++ b/notebooks/v2/create_instance_single_user.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_create_instance_single_user] import sys -# [START notebooks_v2_create_instance_single_user] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -38,9 +37,7 @@ def create_instance_single_user( """ client = notebooks_v2.NotebookServiceClient() parent = f"projects/{project_id}/locations/{location_id}" - service_account_email = ( - f"{service_account_id}@{project_id}.iam.gserviceaccount.com" - ) + service_account_email = f"{service_account_id}@{project_id}.iam.gserviceaccount.com" service_account = notebooks_v2.ServiceAccount(email=service_account_email) gce_setup = notebooks_v2.GceSetup( machine_type="e2-standard-4", @@ -72,42 +69,3 @@ def create_instance_single_user( # [END notebooks_v2_create_instance_single_user] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Create a single-user Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-single-user-instance", - help="The ID for the new Workbench instance.", - ) - parser.add_argument( - "--instance-owner", - default="example-user@example.com", - help="The email address of the instance owner.", - ) - parser.add_argument( - "--service-account-id", - default="example-sa", - help="The ID of the custom VM service account.", - ) - args = parser.parse_args() - create_instance_single_user( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - instance_owner=args.instance_owner, - service_account_id=args.service_account_id, - ) diff --git a/notebooks/v2/create_instance_single_user_test.py b/notebooks/v2/create_instance_single_user_test.py new file mode 100644 index 00000000000..9d4b19342ec --- /dev/null +++ b/notebooks/v2/create_instance_single_user_test.py @@ -0,0 +1,74 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import create_instance_single_user + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-single-user-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_single_user_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + creator="example-user@example.com", + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup(machine_type="e2-standard-4"), + ) + mock_client.create_instance.return_value = operation + + create_instance_single_user.create_instance_single_user( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Creator: example-user@example.com" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_single_user_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance_single_user.create_instance_single_user( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/create_instance_test.py b/notebooks/v2/create_instance_test.py new file mode 100644 index 00000000000..646e3def18e --- /dev/null +++ b/notebooks/v2/create_instance_test.py @@ -0,0 +1,71 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import create_instance + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + proxy_uri="example-proxy.notebooks.googleusercontent.com", + gce_setup=notebooks_v2.GceSetup(machine_type="e2-standard-4"), + ) + mock_client.create_instance.return_value = operation + + create_instance.create_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Machine Type: e2-standard-4" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_create_instance_already_exists( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.create_instance.side_effect = exceptions.AlreadyExists( + "Instance already exists" + ) + + create_instance.create_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/delete_instance.py b/notebooks/v2/delete_instance.py index 95e33b95c25..5f5a7724561 100644 --- a/notebooks/v2/delete_instance.py +++ b/notebooks/v2/delete_instance.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_delete_instance] import sys -# [START notebooks_v2_delete_instance] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -42,7 +41,7 @@ def delete_instance( operation.result() print(f"Resource ID: {name}") print(f"Location: {location_id}") - print(f"Status: Deleted") + print(f"Operation Done: {operation.done()}") except exceptions.NotFound as err: print( f"error: Instance '{instance_id}' not found in location " @@ -52,30 +51,3 @@ def delete_instance( # [END notebooks_v2_delete_instance] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Delete a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - delete_instance( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/delete_instance_test.py b/notebooks/v2/delete_instance_test.py new file mode 100644 index 00000000000..bd72ca4fc37 --- /dev/null +++ b/notebooks/v2/delete_instance_test.py @@ -0,0 +1,63 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +import pytest + +import delete_instance + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_delete_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = None + operation.done.return_value = True + mock_client.delete_instance.return_value = operation + + delete_instance.delete_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Operation Done: True" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_delete_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.delete_instance.side_effect = exceptions.NotFound("Instance not found") + + delete_instance.delete_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/diagnose_instance.py b/notebooks/v2/diagnose_instance.py index 715f541c612..2a6b33cdb71 100644 --- a/notebooks/v2/diagnose_instance.py +++ b/notebooks/v2/diagnose_instance.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_diagnose_instance] import sys -# [START notebooks_v2_diagnose_instance] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -62,36 +61,3 @@ def diagnose_instance( # [END notebooks_v2_diagnose_instance] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Run diagnostics on a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - parser.add_argument( - "--gcs-bucket", - default="gs://example-bucket", - help="The Cloud Storage bucket URI for diagnostic output.", - ) - args = parser.parse_args() - diagnose_instance( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - gcs_bucket=args.gcs_bucket, - ) diff --git a/notebooks/v2/diagnose_instance_test.py b/notebooks/v2/diagnose_instance_test.py new file mode 100644 index 00000000000..b20c25ef6a0 --- /dev/null +++ b/notebooks/v2/diagnose_instance_test.py @@ -0,0 +1,69 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import diagnose_instance + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_diagnose_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + health_state=notebooks_v2.HealthState.HEALTHY, + state=notebooks_v2.State.ACTIVE, + ) + mock_client.diagnose_instance.return_value = operation + + diagnose_instance.diagnose_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Health State: HEALTHY" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_diagnose_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.diagnose_instance.side_effect = exceptions.NotFound( + "Instance not found" + ) + + diagnose_instance.diagnose_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/get_instance.py b/notebooks/v2/get_instance.py index c66a4d417bc..f09c36d1b3d 100644 --- a/notebooks/v2/get_instance.py +++ b/notebooks/v2/get_instance.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_get_instance] import sys -# [START notebooks_v2_get_instance] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -52,30 +51,3 @@ def get_instance( # [END notebooks_v2_get_instance] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Get details for a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - get_instance( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/get_instance_iam_policy.py b/notebooks/v2/get_instance_iam_policy.py index 2515e9fdfe4..8c8d3660759 100644 --- a/notebooks/v2/get_instance_iam_policy.py +++ b/notebooks/v2/get_instance_iam_policy.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_get_instance_iam_policy] import sys -# [START notebooks_v2_get_instance_iam_policy] from google.api_core import exceptions from google.cloud import notebooks_v2 from google.iam.v1 import iam_policy_pb2 @@ -53,30 +52,3 @@ def get_instance_iam_policy( # [END notebooks_v2_get_instance_iam_policy] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Get the IAM policy for a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - get_instance_iam_policy( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/get_instance_iam_policy_test.py b/notebooks/v2/get_instance_iam_policy_test.py new file mode 100644 index 00000000000..90aa91d65af --- /dev/null +++ b/notebooks/v2/get_instance_iam_policy_test.py @@ -0,0 +1,70 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.iam.v1 import policy_pb2 +import pytest + +import get_instance_iam_policy + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_get_instance_iam_policy_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_iam_policy.return_value = policy_pb2.Policy( + version=1, + etag=b"etag-1", + bindings=[ + policy_pb2.Binding( + role="roles/notebooks.runner", + members=["user:example-user@example.com"], + ) + ], + ) + + get_instance_iam_policy.get_instance_iam_policy( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Role: roles/notebooks.runner" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_get_instance_iam_policy_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_iam_policy.side_effect = exceptions.NotFound("Instance not found") + + get_instance_iam_policy.get_instance_iam_policy( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/get_instance_test.py b/notebooks/v2/get_instance_test.py new file mode 100644 index 00000000000..5e8b01baadd --- /dev/null +++ b/notebooks/v2/get_instance_test.py @@ -0,0 +1,66 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import get_instance + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_get_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_instance.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + id="1234567890", + creator="example-user@example.com", + proxy_uri="example-proxy.notebooks.googleusercontent.com", + state=notebooks_v2.State.ACTIVE, + health_state=notebooks_v2.HealthState.HEALTHY, + gce_setup=notebooks_v2.GceSetup(machine_type="e2-standard-4"), + ) + + get_instance.get_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Health State: HEALTHY" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_get_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_instance.side_effect = exceptions.NotFound("Instance not found") + + get_instance.get_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/instances_test.py b/notebooks/v2/instances_test.py deleted file mode 100644 index 1b52edd5acf..00000000000 --- a/notebooks/v2/instances_test.py +++ /dev/null @@ -1,343 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed 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. - -from unittest import mock - -from google.api_core import exceptions -from google.cloud import notebooks_v2 -from google.iam.v1 import iam_policy_pb2 -from google.iam.v1 import policy_pb2 -import pytest - -import check_instance_upgradability -import create_instance -import create_instance_cmek -import create_instance_confidential -import create_instance_custom_container -import create_instance_gpu -import create_instance_hardened -import create_instance_private_vpc -import create_instance_single_user -import delete_instance -import diagnose_instance -import get_instance -import get_instance_iam_policy -import list_instances -import reset_instance -import rollback_instance -import set_instance_iam_policy -import start_instance -import stop_instance -import test_instance_iam_permissions -import update_instance_machine_type -import update_instance_metadata -import update_instance_shielded_config -import upgrade_instance - -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" -INSTANCE_ID = "example-instance" -INSTANCE_NAME = ( - f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -) - - -def _make_mock_instance( - name: str = INSTANCE_NAME, - machine_type: str = "e2-standard-4", -) -> notebooks_v2.Instance: - return notebooks_v2.Instance( - name=name, - id="1234567890", - creator="example-user@example.com", - proxy_uri="example-proxy.notebooks.googleusercontent.com", - state=notebooks_v2.State.ACTIVE, - health_state=notebooks_v2.HealthState.HEALTHY, - gce_setup=notebooks_v2.GceSetup( - machine_type=machine_type, - disable_public_ip=True, - metadata={"idle-timeout-seconds": "3600"}, - boot_disk=notebooks_v2.BootDisk( - disk_size_gb=200, - disk_encryption=notebooks_v2.DiskEncryption.CMEK, - kms_key=( - f"projects/{PROJECT_ID}/locations/us-central1/" - "keyRings/example-key-ring/cryptoKeys/example-key" - ), - ), - container_image=notebooks_v2.ContainerImage( - repository="gcr.io/deeplearning-platform-release/base-cpu", - tag="latest", - ), - shielded_instance_config=notebooks_v2.ShieldedInstanceConfig( - enable_secure_boot=True, - enable_vtpm=True, - enable_integrity_monitoring=True, - ), - ), - ) - - -def _make_mock_lro(result_obj: object = None) -> mock.MagicMock: - operation = mock.MagicMock() - operation.operation.name = ( - f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" - ) - operation.result.return_value = result_obj - return operation - - -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") -def test_list_instances_success( - mock_client_cls: mock.MagicMock, - capsys: pytest.CaptureFixture[str], -) -> None: - mock_client = mock_client_cls.return_value - mock_client.list_instances.return_value = [_make_mock_instance()] - - list_instances.list_instances(PROJECT_ID, LOCATION_ID) - - captured = capsys.readouterr() - assert f"Resource ID: {INSTANCE_NAME}" in captured.out - assert "State: ACTIVE" in captured.out - assert captured.err == "" - - -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") -def test_list_instances_not_found( - mock_client_cls: mock.MagicMock, - capsys: pytest.CaptureFixture[str], -) -> None: - mock_client = mock_client_cls.return_value - mock_client.list_instances.side_effect = exceptions.NotFound("Zone not found") - - list_instances.list_instances(PROJECT_ID, LOCATION_ID) - - captured = capsys.readouterr() - assert "error:" in captured.err - assert LOCATION_ID in captured.err - - -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") -def test_get_instance_success( - mock_client_cls: mock.MagicMock, - capsys: pytest.CaptureFixture[str], -) -> None: - mock_client = mock_client_cls.return_value - mock_client.get_instance.return_value = _make_mock_instance() - - get_instance.get_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - - captured = capsys.readouterr() - assert f"Resource ID: {INSTANCE_NAME}" in captured.out - assert "Health State: HEALTHY" in captured.out - assert "State: ACTIVE" in captured.out - - -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") -def test_get_instance_not_found( - mock_client_cls: mock.MagicMock, - capsys: pytest.CaptureFixture[str], -) -> None: - mock_client = mock_client_cls.return_value - mock_client.get_instance.side_effect = exceptions.NotFound("Instance missing") - - get_instance.get_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - - captured = capsys.readouterr() - assert "error:" in captured.err - assert INSTANCE_ID in captured.err - - -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") -def test_create_instance_success( - mock_client_cls: mock.MagicMock, - capsys: pytest.CaptureFixture[str], -) -> None: - mock_client = mock_client_cls.return_value - mock_client.create_instance.return_value = _make_mock_lro( - _make_mock_instance() - ) - - create_instance.create_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - - captured = capsys.readouterr() - assert "Waiting for operation to complete:" in captured.out - assert f"Resource ID: {INSTANCE_NAME}" in captured.out - assert "Machine Type: e2-standard-4" in captured.out - - -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") -def test_create_instance_already_exists( - mock_client_cls: mock.MagicMock, - capsys: pytest.CaptureFixture[str], -) -> None: - mock_client = mock_client_cls.return_value - mock_client.create_instance.side_effect = exceptions.AlreadyExists( - "Duplicate instance" - ) - - create_instance.create_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - - captured = capsys.readouterr() - assert "error:" in captured.err - assert INSTANCE_ID in captured.err - - -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") -def test_stop_start_reset_delete_instance( - mock_client_cls: mock.MagicMock, - capsys: pytest.CaptureFixture[str], -) -> None: - mock_client = mock_client_cls.return_value - mock_client.stop_instance.return_value = _make_mock_lro( - _make_mock_instance() - ) - mock_client.start_instance.return_value = _make_mock_lro( - _make_mock_instance() - ) - mock_client.reset_instance.return_value = _make_mock_lro( - _make_mock_instance() - ) - mock_client.delete_instance.return_value = _make_mock_lro(None) - - stop_instance.stop_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - start_instance.start_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - reset_instance.reset_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - delete_instance.delete_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - - captured = capsys.readouterr() - assert f"Resource ID: {INSTANCE_NAME}" in captured.out - assert "Status: Deleted" in captured.out - - -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") -def test_update_and_maintenance_operations( - mock_client_cls: mock.MagicMock, - capsys: pytest.CaptureFixture[str], -) -> None: - mock_client = mock_client_cls.return_value - mock_client.update_instance.return_value = _make_mock_lro( - _make_mock_instance(machine_type="n1-standard-8") - ) - mock_client.check_instance_upgradability.return_value = ( - notebooks_v2.CheckInstanceUpgradabilityResponse( - upgradeable=True, - upgrade_version="m120", - upgrade_info="Upgrade available", - upgrade_image="gcr.io/deeplearning-platform-release/base-cpu:m120", - ) - ) - mock_client.upgrade_instance.return_value = _make_mock_lro( - _make_mock_instance() - ) - mock_client.rollback_instance.return_value = _make_mock_lro( - _make_mock_instance() - ) - mock_client.diagnose_instance.return_value = _make_mock_lro( - _make_mock_instance() - ) - - update_instance_machine_type.update_instance_machine_type( - PROJECT_ID, LOCATION_ID, INSTANCE_ID, "n1-standard-8" - ) - update_instance_metadata.update_instance_metadata( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - update_instance_shielded_config.update_instance_shielded_config( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - check_instance_upgradability.check_instance_upgradability( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - upgrade_instance.upgrade_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - rollback_instance.rollback_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - diagnose_instance.diagnose_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) - - captured = capsys.readouterr() - assert "Machine Type: n1-standard-8" in captured.out - assert "Secure Boot Enabled: True" in captured.out - assert "Upgradeable: True" in captured.out - assert "Health State: HEALTHY" in captured.out - - -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") -def test_specialized_create_and_iam_operations( - mock_client_cls: mock.MagicMock, - capsys: pytest.CaptureFixture[str], -) -> None: - mock_client = mock_client_cls.return_value - mock_client.create_instance.return_value = _make_mock_lro( - _make_mock_instance() - ) - mock_policy = policy_pb2.Policy( - version=1, - etag=b"etag-1", - bindings=[ - policy_pb2.Binding( - role="roles/notebooks.runner", - members=["user:example-user@example.com"], - ) - ], - ) - mock_client.get_iam_policy.return_value = mock_policy - mock_client.set_iam_policy.return_value = mock_policy - mock_client.test_iam_permissions.return_value = ( - iam_policy_pb2.TestIamPermissionsResponse( - permissions=["notebooks.instances.get"] - ) - ) - - create_instance_gpu.create_instance_gpu( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - create_instance_custom_container.create_instance_custom_container( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - create_instance_cmek.create_instance_cmek( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - create_instance_confidential.create_instance_confidential( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - create_instance_private_vpc.create_instance_private_vpc( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - create_instance_single_user.create_instance_single_user( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - create_instance_hardened.create_instance_hardened( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - get_instance_iam_policy.get_instance_iam_policy( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - set_instance_iam_policy.set_instance_iam_policy( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - test_instance_iam_permissions.test_instance_iam_permissions( - PROJECT_ID, LOCATION_ID, INSTANCE_ID - ) - - captured = capsys.readouterr() - assert ( - "Container Repository: gcr.io/deeplearning-platform-release/base-cpu" - in captured.out - ) - assert "Boot Disk Encryption: CMEK" in captured.out - assert "Disable Public IP: True" in captured.out - assert "Creator: example-user@example.com" in captured.out - assert "Secure Boot Enabled: True" in captured.out - assert "Role: roles/notebooks.runner" in captured.out - assert "Allowed Permissions: ['notebooks.instances.get']" in captured.out diff --git a/notebooks/v2/list_instances.py b/notebooks/v2/list_instances.py index 84f11113160..d00ff9ba1ec 100644 --- a/notebooks/v2/list_instances.py +++ b/notebooks/v2/list_instances.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_list_instances] import sys -# [START notebooks_v2_list_instances] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -51,24 +50,3 @@ def list_instances( # [END notebooks_v2_list_instances] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="List Workbench instances." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - args = parser.parse_args() - list_instances( - project_id=args.project_id, - location_id=args.location_id, - ) diff --git a/notebooks/v2/list_instances_test.py b/notebooks/v2/list_instances_test.py new file mode 100644 index 00000000000..c99b9520d05 --- /dev/null +++ b/notebooks/v2/list_instances_test.py @@ -0,0 +1,65 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import list_instances + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_list_instances_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.list_instances.return_value = [ + notebooks_v2.Instance( + name=INSTANCE_NAME, + id="1234567890", + creator="example-user@example.com", + proxy_uri="example-proxy.notebooks.googleusercontent.com", + state=notebooks_v2.State.ACTIVE, + ) + ] + + list_instances.list_instances(PROJECT_ID, LOCATION_ID) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_list_instances_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.list_instances.side_effect = exceptions.NotFound("Zone not found") + + list_instances.list_instances(PROJECT_ID, LOCATION_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert LOCATION_ID in captured.err diff --git a/notebooks/v2/noxfile_config.py b/notebooks/v2/noxfile_config.py index 719a62bc276..cb3ff7a91a6 100644 --- a/notebooks/v2/noxfile_config.py +++ b/notebooks/v2/noxfile_config.py @@ -4,7 +4,7 @@ # 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 +# 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, @@ -16,7 +16,7 @@ TEST_CONFIG_OVERRIDE = { # You can opt out from the test for specific Python versions. - "ignored_versions": [], + "ignored_versions": ["3.8", "3.9", "3.11", "3.12", "3.13"], # Old samples are opted out of enforcing Python type hints # All new samples should feature them "enforce_type_hints": True, diff --git a/notebooks/v2/reset_instance.py b/notebooks/v2/reset_instance.py index 2c95aa60be8..b24c88f2f37 100644 --- a/notebooks/v2/reset_instance.py +++ b/notebooks/v2/reset_instance.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_reset_instance] import sys -# [START notebooks_v2_reset_instance] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -51,30 +50,3 @@ def reset_instance( # [END notebooks_v2_reset_instance] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Reset a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - reset_instance( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/reset_instance_test.py b/notebooks/v2/reset_instance_test.py new file mode 100644 index 00000000000..f6d141f1935 --- /dev/null +++ b/notebooks/v2/reset_instance_test.py @@ -0,0 +1,66 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import reset_instance + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_reset_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + ) + mock_client.reset_instance.return_value = operation + + reset_instance.reset_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_reset_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.reset_instance.side_effect = exceptions.NotFound("Instance not found") + + reset_instance.reset_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/rollback_instance.py b/notebooks/v2/rollback_instance.py index 696ef80dd9c..eed58adde3e 100644 --- a/notebooks/v2/rollback_instance.py +++ b/notebooks/v2/rollback_instance.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_rollback_instance] import sys -# [START notebooks_v2_rollback_instance] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -60,42 +59,3 @@ def rollback_instance( # [END notebooks_v2_rollback_instance] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Roll back a Workbench instance to a snapshot." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - parser.add_argument( - "--snapshot-id", - default="example-snapshot", - help="The ID of the target Compute Engine snapshot.", - ) - parser.add_argument( - "--revision-id", - default="1", - help="The revision ID for rollback.", - ) - args = parser.parse_args() - rollback_instance( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - snapshot_id=args.snapshot_id, - revision_id=args.revision_id, - ) diff --git a/notebooks/v2/rollback_instance_test.py b/notebooks/v2/rollback_instance_test.py new file mode 100644 index 00000000000..25176a19b22 --- /dev/null +++ b/notebooks/v2/rollback_instance_test.py @@ -0,0 +1,67 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import rollback_instance + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_rollback_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + ) + mock_client.rollback_instance.return_value = operation + + rollback_instance.rollback_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_rollback_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.rollback_instance.side_effect = exceptions.NotFound( + "Instance not found" + ) + + rollback_instance.rollback_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/set_instance_iam_policy.py b/notebooks/v2/set_instance_iam_policy.py index f1471aceb5f..998278bfb4f 100644 --- a/notebooks/v2/set_instance_iam_policy.py +++ b/notebooks/v2/set_instance_iam_policy.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_set_instance_iam_policy] import sys -# [START notebooks_v2_set_instance_iam_policy] from google.api_core import exceptions from google.cloud import notebooks_v2 from google.iam.v1 import iam_policy_pb2 @@ -29,7 +28,10 @@ def set_instance_iam_policy( member: str = "user:example-user@example.com", role: str = "roles/notebooks.runner", ) -> None: - """Sets the IAM policy for a Workbench instance. + """Sets the IAM policy for a Workbench instance using read-modify-write. + + Retrieves the existing IAM policy first so that appending the new role + binding preserves existing access controls and optimistic concurrency etags. Args: project_id: The Google Cloud project ID. @@ -40,23 +42,25 @@ def set_instance_iam_policy( """ client = notebooks_v2.NotebookServiceClient() name = f"projects/{project_id}/locations/{location_id}/instances/{instance_id}" + get_request = iam_policy_pb2.GetIamPolicyRequest(resource=name) binding = policy_pb2.Binding( role=role, members=[member], ) - policy = policy_pb2.Policy(bindings=[binding]) - request = iam_policy_pb2.SetIamPolicyRequest( - resource=name, - policy=policy, - ) try: - response = client.set_iam_policy(request=request) + # Read the current IAM policy first to preserve existing bindings and etag. + policy = client.get_iam_policy(request=get_request) + policy.bindings.append(binding) + set_request = iam_policy_pb2.SetIamPolicyRequest( + resource=name, + policy=policy, + ) + response = client.set_iam_policy(request=set_request) print(f"Resource ID: {name}") + print(f"Bindings Count: {len(response.bindings)}") print(f"Policy Etag: {response.etag}") print(f"Policy Version: {response.version}") - for item in response.bindings: - print(f"Role: {item.role}") except exceptions.NotFound as err: print( f"error: Instance '{instance_id}' not found in location " @@ -66,42 +70,3 @@ def set_instance_iam_policy( # [END notebooks_v2_set_instance_iam_policy] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Set the IAM policy for a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - parser.add_argument( - "--member", - default="user:example-user@example.com", - help="The IAM principal to grant access to.", - ) - parser.add_argument( - "--role", - default="roles/notebooks.runner", - help="The IAM role to grant.", - ) - args = parser.parse_args() - set_instance_iam_policy( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - member=args.member, - role=args.role, - ) diff --git a/notebooks/v2/set_instance_iam_policy_test.py b/notebooks/v2/set_instance_iam_policy_test.py new file mode 100644 index 00000000000..f593184e997 --- /dev/null +++ b/notebooks/v2/set_instance_iam_policy_test.py @@ -0,0 +1,79 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.iam.v1 import policy_pb2 +import pytest + +import set_instance_iam_policy + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_set_instance_iam_policy_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + existing_policy = policy_pb2.Policy( + version=1, + etag=b"etag-1", + bindings=[], + ) + updated_policy = policy_pb2.Policy( + version=1, + etag=b"etag-2", + bindings=[ + policy_pb2.Binding( + role="roles/notebooks.runner", + members=["user:example-user@example.com"], + ) + ], + ) + mock_client.get_iam_policy.return_value = existing_policy + mock_client.set_iam_policy.return_value = updated_policy + + set_instance_iam_policy.set_instance_iam_policy( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + mock_client.get_iam_policy.assert_called_once() + mock_client.set_iam_policy.assert_called_once() + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Bindings Count: 1" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_set_instance_iam_policy_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.get_iam_policy.side_effect = exceptions.NotFound("Instance not found") + + set_instance_iam_policy.set_instance_iam_policy( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/start_instance.py b/notebooks/v2/start_instance.py index 84ae2c4f868..0a4ef2f64d3 100644 --- a/notebooks/v2/start_instance.py +++ b/notebooks/v2/start_instance.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_start_instance] import sys -# [START notebooks_v2_start_instance] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -52,30 +51,3 @@ def start_instance( # [END notebooks_v2_start_instance] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Start a stopped Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - start_instance( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/start_instance_test.py b/notebooks/v2/start_instance_test.py new file mode 100644 index 00000000000..b9525e7a24d --- /dev/null +++ b/notebooks/v2/start_instance_test.py @@ -0,0 +1,67 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import start_instance + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_start_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + proxy_uri="example-proxy.notebooks.googleusercontent.com", + state=notebooks_v2.State.ACTIVE, + ) + mock_client.start_instance.return_value = operation + + start_instance.start_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_start_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.start_instance.side_effect = exceptions.NotFound("Instance not found") + + start_instance.start_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/stop_instance.py b/notebooks/v2/stop_instance.py index 2b0d744b091..1bac8a3bffd 100644 --- a/notebooks/v2/stop_instance.py +++ b/notebooks/v2/stop_instance.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_stop_instance] import sys -# [START notebooks_v2_stop_instance] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -51,30 +50,3 @@ def stop_instance( # [END notebooks_v2_stop_instance] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Stop a running Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - stop_instance( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/stop_instance_test.py b/notebooks/v2/stop_instance_test.py new file mode 100644 index 00000000000..47bf491676c --- /dev/null +++ b/notebooks/v2/stop_instance_test.py @@ -0,0 +1,66 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import stop_instance + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_stop_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.STOPPED, + ) + mock_client.stop_instance.return_value = operation + + stop_instance.stop_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "Waiting for operation to complete:" in captured.out + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: STOPPED" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_stop_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.stop_instance.side_effect = exceptions.NotFound("Instance not found") + + stop_instance.stop_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/update_instance_machine_type.py b/notebooks/v2/update_instance_machine_type.py index 59f34293551..b133cde593a 100644 --- a/notebooks/v2/update_instance_machine_type.py +++ b/notebooks/v2/update_instance_machine_type.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_update_instance_machine_type] import sys -# [START notebooks_v2_update_instance_machine_type] from google.api_core import exceptions from google.cloud import notebooks_v2 from google.protobuf import field_mask_pb2 @@ -61,36 +60,3 @@ def update_instance_machine_type( # [END notebooks_v2_update_instance_machine_type] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Update the machine type of a stopped Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - parser.add_argument( - "--machine-type", - default="n1-standard-8", - help="The Compute Engine machine type.", - ) - args = parser.parse_args() - update_instance_machine_type( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - machine_type=args.machine_type, - ) diff --git a/notebooks/v2/update_instance_machine_type_test.py b/notebooks/v2/update_instance_machine_type_test.py new file mode 100644 index 00000000000..6354a8f67fa --- /dev/null +++ b/notebooks/v2/update_instance_machine_type_test.py @@ -0,0 +1,71 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import update_instance_machine_type + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_update_instance_machine_type_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.STOPPED, + gce_setup=notebooks_v2.GceSetup(machine_type="n1-standard-8"), + ) + mock_client.update_instance.return_value = operation + + update_instance_machine_type.update_instance_machine_type( + PROJECT_ID, LOCATION_ID, INSTANCE_ID, "n1-standard-8" + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Machine Type: n1-standard-8" in captured.out + assert "State: STOPPED" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_update_instance_machine_type_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.update_instance.side_effect = exceptions.NotFound("Instance not found") + + update_instance_machine_type.update_instance_machine_type( + PROJECT_ID, LOCATION_ID, INSTANCE_ID, "n1-standard-8" + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/update_instance_metadata.py b/notebooks/v2/update_instance_metadata.py index f0f6c19706b..36067ff3066 100644 --- a/notebooks/v2/update_instance_metadata.py +++ b/notebooks/v2/update_instance_metadata.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_update_instance_metadata] import sys -# [START notebooks_v2_update_instance_metadata] from google.api_core import exceptions from google.cloud import notebooks_v2 from google.protobuf import field_mask_pb2 @@ -61,30 +60,3 @@ def update_instance_metadata( # [END notebooks_v2_update_instance_metadata] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Update custom metadata on a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - update_instance_metadata( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/update_instance_metadata_test.py b/notebooks/v2/update_instance_metadata_test.py new file mode 100644 index 00000000000..d008f0d01d1 --- /dev/null +++ b/notebooks/v2/update_instance_metadata_test.py @@ -0,0 +1,71 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import update_instance_metadata + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_update_instance_metadata_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + gce_setup=notebooks_v2.GceSetup(metadata={"idle-timeout-seconds": "3600"}), + ) + mock_client.update_instance.return_value = operation + + update_instance_metadata.update_instance_metadata( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "idle-timeout-seconds" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_update_instance_metadata_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.update_instance.side_effect = exceptions.NotFound("Instance not found") + + update_instance_metadata.update_instance_metadata( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/update_instance_shielded_config.py b/notebooks/v2/update_instance_shielded_config.py index f9a45c63306..a22f7d78a7f 100644 --- a/notebooks/v2/update_instance_shielded_config.py +++ b/notebooks/v2/update_instance_shielded_config.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_update_instance_shielded_config] import sys -# [START notebooks_v2_update_instance_shielded_config] from google.api_core import exceptions from google.cloud import notebooks_v2 from google.protobuf import field_mask_pb2 @@ -58,9 +57,7 @@ def update_instance_shielded_config( operation = client.update_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") response = operation.result() - secure_boot = ( - response.gce_setup.shielded_instance_config.enable_secure_boot - ) + secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot print(f"Resource ID: {response.name}") print(f"Secure Boot Enabled: {secure_boot}") print(f"State: {response.state.name}") @@ -73,30 +70,3 @@ def update_instance_shielded_config( # [END notebooks_v2_update_instance_shielded_config] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Update the Shielded VM config of a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - update_instance_shielded_config( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/update_instance_shielded_config_test.py b/notebooks/v2/update_instance_shielded_config_test.py new file mode 100644 index 00000000000..ddb719b48ef --- /dev/null +++ b/notebooks/v2/update_instance_shielded_config_test.py @@ -0,0 +1,77 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import update_instance_shielded_config + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_update_instance_shielded_config_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.STOPPED, + gce_setup=notebooks_v2.GceSetup( + shielded_instance_config=notebooks_v2.ShieldedInstanceConfig( + enable_secure_boot=True, + enable_vtpm=True, + enable_integrity_monitoring=True, + ) + ), + ) + mock_client.update_instance.return_value = operation + + update_instance_shielded_config.update_instance_shielded_config( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Secure Boot Enabled: True" in captured.out + assert "State: STOPPED" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_update_instance_shielded_config_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.update_instance.side_effect = exceptions.NotFound("Instance not found") + + update_instance_shielded_config.update_instance_shielded_config( + PROJECT_ID, LOCATION_ID, INSTANCE_ID + ) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/upgrade_instance.py b/notebooks/v2/upgrade_instance.py index 9e24791c40c..7d4edd93daa 100644 --- a/notebooks/v2/upgrade_instance.py +++ b/notebooks/v2/upgrade_instance.py @@ -12,10 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +# [START notebooks_v2_upgrade_instance] import sys -# [START notebooks_v2_upgrade_instance] from google.api_core import exceptions from google.cloud import notebooks_v2 @@ -51,30 +50,3 @@ def upgrade_instance( # [END notebooks_v2_upgrade_instance] - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Upgrade a Workbench instance." - ) - parser.add_argument( - "--project-id", - required=True, - help="The Google Cloud project ID.", - ) - parser.add_argument( - "--location-id", - default="us-central1-a", - help="The Compute Engine zone.", - ) - parser.add_argument( - "--instance-id", - default="example-instance", - help="The ID of the Workbench instance.", - ) - args = parser.parse_args() - upgrade_instance( - project_id=args.project_id, - location_id=args.location_id, - instance_id=args.instance_id, - ) diff --git a/notebooks/v2/upgrade_instance_test.py b/notebooks/v2/upgrade_instance_test.py new file mode 100644 index 00000000000..8232fd8281e --- /dev/null +++ b/notebooks/v2/upgrade_instance_test.py @@ -0,0 +1,65 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. + +from unittest import mock + +from google.api_core import exceptions +from google.cloud import notebooks_v2 +import pytest + +import upgrade_instance + +PROJECT_ID = "example-project" +LOCATION_ID = "us-central1-a" +INSTANCE_ID = "example-instance" +INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_upgrade_instance_success( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + operation = mock.MagicMock() + operation.operation.name = ( + f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/operations/op-123" + ) + operation.result.return_value = notebooks_v2.Instance( + name=INSTANCE_NAME, + state=notebooks_v2.State.ACTIVE, + ) + mock_client.upgrade_instance.return_value = operation + + upgrade_instance.upgrade_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "State: ACTIVE" in captured.out + assert captured.err == "" + + +@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +def test_upgrade_instance_not_found( + mock_client_cls: mock.MagicMock, + capsys: pytest.CaptureFixture[str], +) -> None: + mock_client = mock_client_cls.return_value + mock_client.upgrade_instance.side_effect = exceptions.NotFound("Instance not found") + + upgrade_instance.upgrade_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + + captured = capsys.readouterr() + assert "error:" in captured.err + assert INSTANCE_ID in captured.err From 9341a1d978b1c47ba143b6a942d754441d8e195d Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Tue, 3 Nov 2026 19:10:17 +0000 Subject: [PATCH 3/8] refactor(notebooks): move v2 per-sample unit tests into notebooks/v2/tests/ Problem: 1. Placing all 24 per-sample unit test files directly in notebooks/v2/ cluttered the top-level sample directory. 2. When GitHub Actions lint-changed.sh runs flake8 from the repository root, top-level bare module imports in notebooks/v2/ were not recognized as local application imports by flake8-import-order. Solution: 1. Move all 24 _test.py files into notebooks/v2/tests/ and add package __init__.py files. 2. Use relative imports (from .. import ) in notebooks/v2/tests/*_test.py so flake8-import-order and pytest succeed whether invoked from the repository root or from notebooks/v2/. --- notebooks/v2/__init__.py | 13 +++++++++++++ notebooks/v2/tests/__init__.py | 13 +++++++++++++ .../check_instance_iam_permissions_test.py | 2 +- .../check_instance_upgradability_test.py | 2 +- .../v2/{ => tests}/create_instance_cmek_test.py | 2 +- .../create_instance_confidential_test.py | 2 +- .../create_instance_custom_container_test.py | 2 +- .../v2/{ => tests}/create_instance_gpu_test.py | 2 +- .../v2/{ => tests}/create_instance_hardened_test.py | 2 +- .../{ => tests}/create_instance_private_vpc_test.py | 2 +- .../{ => tests}/create_instance_single_user_test.py | 2 +- notebooks/v2/{ => tests}/create_instance_test.py | 2 +- notebooks/v2/{ => tests}/delete_instance_test.py | 2 +- notebooks/v2/{ => tests}/diagnose_instance_test.py | 2 +- .../v2/{ => tests}/get_instance_iam_policy_test.py | 2 +- notebooks/v2/{ => tests}/get_instance_test.py | 2 +- notebooks/v2/{ => tests}/list_instances_test.py | 2 +- notebooks/v2/{ => tests}/reset_instance_test.py | 2 +- notebooks/v2/{ => tests}/rollback_instance_test.py | 2 +- .../v2/{ => tests}/set_instance_iam_policy_test.py | 2 +- notebooks/v2/{ => tests}/start_instance_test.py | 2 +- notebooks/v2/{ => tests}/stop_instance_test.py | 2 +- .../update_instance_machine_type_test.py | 2 +- .../v2/{ => tests}/update_instance_metadata_test.py | 2 +- .../update_instance_shielded_config_test.py | 2 +- notebooks/v2/{ => tests}/upgrade_instance_test.py | 2 +- 26 files changed, 50 insertions(+), 24 deletions(-) create mode 100644 notebooks/v2/__init__.py create mode 100644 notebooks/v2/tests/__init__.py rename notebooks/v2/{ => tests}/check_instance_iam_permissions_test.py (97%) rename notebooks/v2/{ => tests}/check_instance_upgradability_test.py (98%) rename notebooks/v2/{ => tests}/create_instance_cmek_test.py (98%) rename notebooks/v2/{ => tests}/create_instance_confidential_test.py (98%) rename notebooks/v2/{ => tests}/create_instance_custom_container_test.py (98%) rename notebooks/v2/{ => tests}/create_instance_gpu_test.py (98%) rename notebooks/v2/{ => tests}/create_instance_hardened_test.py (98%) rename notebooks/v2/{ => tests}/create_instance_private_vpc_test.py (98%) rename notebooks/v2/{ => tests}/create_instance_single_user_test.py (98%) rename notebooks/v2/{ => tests}/create_instance_test.py (98%) rename notebooks/v2/{ => tests}/delete_instance_test.py (98%) rename notebooks/v2/{ => tests}/diagnose_instance_test.py (98%) rename notebooks/v2/{ => tests}/get_instance_iam_policy_test.py (98%) rename notebooks/v2/{ => tests}/get_instance_test.py (98%) rename notebooks/v2/{ => tests}/list_instances_test.py (98%) rename notebooks/v2/{ => tests}/reset_instance_test.py (98%) rename notebooks/v2/{ => tests}/rollback_instance_test.py (98%) rename notebooks/v2/{ => tests}/set_instance_iam_policy_test.py (98%) rename notebooks/v2/{ => tests}/start_instance_test.py (98%) rename notebooks/v2/{ => tests}/stop_instance_test.py (98%) rename notebooks/v2/{ => tests}/update_instance_machine_type_test.py (98%) rename notebooks/v2/{ => tests}/update_instance_metadata_test.py (98%) rename notebooks/v2/{ => tests}/update_instance_shielded_config_test.py (98%) rename notebooks/v2/{ => tests}/upgrade_instance_test.py (98%) diff --git a/notebooks/v2/__init__.py b/notebooks/v2/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/notebooks/v2/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. diff --git a/notebooks/v2/tests/__init__.py b/notebooks/v2/tests/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/notebooks/v2/tests/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed 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. diff --git a/notebooks/v2/check_instance_iam_permissions_test.py b/notebooks/v2/tests/check_instance_iam_permissions_test.py similarity index 97% rename from notebooks/v2/check_instance_iam_permissions_test.py rename to notebooks/v2/tests/check_instance_iam_permissions_test.py index d3b0af85fa7..c15621d946a 100644 --- a/notebooks/v2/check_instance_iam_permissions_test.py +++ b/notebooks/v2/tests/check_instance_iam_permissions_test.py @@ -18,7 +18,7 @@ from google.iam.v1 import iam_policy_pb2 import pytest -import check_instance_iam_permissions +from .. import check_instance_iam_permissions PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/check_instance_upgradability_test.py b/notebooks/v2/tests/check_instance_upgradability_test.py similarity index 98% rename from notebooks/v2/check_instance_upgradability_test.py rename to notebooks/v2/tests/check_instance_upgradability_test.py index 672336c868c..8c81ed1d6cc 100644 --- a/notebooks/v2/check_instance_upgradability_test.py +++ b/notebooks/v2/tests/check_instance_upgradability_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import check_instance_upgradability +from .. import check_instance_upgradability PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/create_instance_cmek_test.py b/notebooks/v2/tests/create_instance_cmek_test.py similarity index 98% rename from notebooks/v2/create_instance_cmek_test.py rename to notebooks/v2/tests/create_instance_cmek_test.py index cfc4e28d874..ccd9cffc52e 100644 --- a/notebooks/v2/create_instance_cmek_test.py +++ b/notebooks/v2/tests/create_instance_cmek_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import create_instance_cmek +from .. import create_instance_cmek PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/create_instance_confidential_test.py b/notebooks/v2/tests/create_instance_confidential_test.py similarity index 98% rename from notebooks/v2/create_instance_confidential_test.py rename to notebooks/v2/tests/create_instance_confidential_test.py index f48d0092931..59fae743b2a 100644 --- a/notebooks/v2/create_instance_confidential_test.py +++ b/notebooks/v2/tests/create_instance_confidential_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import create_instance_confidential +from .. import create_instance_confidential PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/create_instance_custom_container_test.py b/notebooks/v2/tests/create_instance_custom_container_test.py similarity index 98% rename from notebooks/v2/create_instance_custom_container_test.py rename to notebooks/v2/tests/create_instance_custom_container_test.py index cda0484c0bb..473d86c794d 100644 --- a/notebooks/v2/create_instance_custom_container_test.py +++ b/notebooks/v2/tests/create_instance_custom_container_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import create_instance_custom_container +from .. import create_instance_custom_container PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/create_instance_gpu_test.py b/notebooks/v2/tests/create_instance_gpu_test.py similarity index 98% rename from notebooks/v2/create_instance_gpu_test.py rename to notebooks/v2/tests/create_instance_gpu_test.py index 1cc5d9de4aa..11e245a0f71 100644 --- a/notebooks/v2/create_instance_gpu_test.py +++ b/notebooks/v2/tests/create_instance_gpu_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import create_instance_gpu +from .. import create_instance_gpu PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/create_instance_hardened_test.py b/notebooks/v2/tests/create_instance_hardened_test.py similarity index 98% rename from notebooks/v2/create_instance_hardened_test.py rename to notebooks/v2/tests/create_instance_hardened_test.py index 8e338a053c4..86f0270df89 100644 --- a/notebooks/v2/create_instance_hardened_test.py +++ b/notebooks/v2/tests/create_instance_hardened_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import create_instance_hardened +from .. import create_instance_hardened PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/create_instance_private_vpc_test.py b/notebooks/v2/tests/create_instance_private_vpc_test.py similarity index 98% rename from notebooks/v2/create_instance_private_vpc_test.py rename to notebooks/v2/tests/create_instance_private_vpc_test.py index ecdaae661df..a31565080f4 100644 --- a/notebooks/v2/create_instance_private_vpc_test.py +++ b/notebooks/v2/tests/create_instance_private_vpc_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import create_instance_private_vpc +from .. import create_instance_private_vpc PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/create_instance_single_user_test.py b/notebooks/v2/tests/create_instance_single_user_test.py similarity index 98% rename from notebooks/v2/create_instance_single_user_test.py rename to notebooks/v2/tests/create_instance_single_user_test.py index 9d4b19342ec..8fbbbf38a2b 100644 --- a/notebooks/v2/create_instance_single_user_test.py +++ b/notebooks/v2/tests/create_instance_single_user_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import create_instance_single_user +from .. import create_instance_single_user PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/create_instance_test.py b/notebooks/v2/tests/create_instance_test.py similarity index 98% rename from notebooks/v2/create_instance_test.py rename to notebooks/v2/tests/create_instance_test.py index 646e3def18e..0ec55a1af3a 100644 --- a/notebooks/v2/create_instance_test.py +++ b/notebooks/v2/tests/create_instance_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import create_instance +from .. import create_instance PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/delete_instance_test.py b/notebooks/v2/tests/delete_instance_test.py similarity index 98% rename from notebooks/v2/delete_instance_test.py rename to notebooks/v2/tests/delete_instance_test.py index bd72ca4fc37..9dad8164089 100644 --- a/notebooks/v2/delete_instance_test.py +++ b/notebooks/v2/tests/delete_instance_test.py @@ -17,7 +17,7 @@ from google.api_core import exceptions import pytest -import delete_instance +from .. import delete_instance PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/diagnose_instance_test.py b/notebooks/v2/tests/diagnose_instance_test.py similarity index 98% rename from notebooks/v2/diagnose_instance_test.py rename to notebooks/v2/tests/diagnose_instance_test.py index b20c25ef6a0..5940be65054 100644 --- a/notebooks/v2/diagnose_instance_test.py +++ b/notebooks/v2/tests/diagnose_instance_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import diagnose_instance +from .. import diagnose_instance PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/get_instance_iam_policy_test.py b/notebooks/v2/tests/get_instance_iam_policy_test.py similarity index 98% rename from notebooks/v2/get_instance_iam_policy_test.py rename to notebooks/v2/tests/get_instance_iam_policy_test.py index 90aa91d65af..3e9e12538cb 100644 --- a/notebooks/v2/get_instance_iam_policy_test.py +++ b/notebooks/v2/tests/get_instance_iam_policy_test.py @@ -18,7 +18,7 @@ from google.iam.v1 import policy_pb2 import pytest -import get_instance_iam_policy +from .. import get_instance_iam_policy PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/get_instance_test.py b/notebooks/v2/tests/get_instance_test.py similarity index 98% rename from notebooks/v2/get_instance_test.py rename to notebooks/v2/tests/get_instance_test.py index 5e8b01baadd..1ecd7e6576a 100644 --- a/notebooks/v2/get_instance_test.py +++ b/notebooks/v2/tests/get_instance_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import get_instance +from .. import get_instance PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/list_instances_test.py b/notebooks/v2/tests/list_instances_test.py similarity index 98% rename from notebooks/v2/list_instances_test.py rename to notebooks/v2/tests/list_instances_test.py index c99b9520d05..00d7872eafc 100644 --- a/notebooks/v2/list_instances_test.py +++ b/notebooks/v2/tests/list_instances_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import list_instances +from .. import list_instances PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/reset_instance_test.py b/notebooks/v2/tests/reset_instance_test.py similarity index 98% rename from notebooks/v2/reset_instance_test.py rename to notebooks/v2/tests/reset_instance_test.py index f6d141f1935..d0fe3cc455d 100644 --- a/notebooks/v2/reset_instance_test.py +++ b/notebooks/v2/tests/reset_instance_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import reset_instance +from .. import reset_instance PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/rollback_instance_test.py b/notebooks/v2/tests/rollback_instance_test.py similarity index 98% rename from notebooks/v2/rollback_instance_test.py rename to notebooks/v2/tests/rollback_instance_test.py index 25176a19b22..18215ba8a18 100644 --- a/notebooks/v2/rollback_instance_test.py +++ b/notebooks/v2/tests/rollback_instance_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import rollback_instance +from .. import rollback_instance PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/set_instance_iam_policy_test.py b/notebooks/v2/tests/set_instance_iam_policy_test.py similarity index 98% rename from notebooks/v2/set_instance_iam_policy_test.py rename to notebooks/v2/tests/set_instance_iam_policy_test.py index f593184e997..bbd35420e36 100644 --- a/notebooks/v2/set_instance_iam_policy_test.py +++ b/notebooks/v2/tests/set_instance_iam_policy_test.py @@ -18,7 +18,7 @@ from google.iam.v1 import policy_pb2 import pytest -import set_instance_iam_policy +from .. import set_instance_iam_policy PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/start_instance_test.py b/notebooks/v2/tests/start_instance_test.py similarity index 98% rename from notebooks/v2/start_instance_test.py rename to notebooks/v2/tests/start_instance_test.py index b9525e7a24d..9ad9158422f 100644 --- a/notebooks/v2/start_instance_test.py +++ b/notebooks/v2/tests/start_instance_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import start_instance +from .. import start_instance PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/stop_instance_test.py b/notebooks/v2/tests/stop_instance_test.py similarity index 98% rename from notebooks/v2/stop_instance_test.py rename to notebooks/v2/tests/stop_instance_test.py index 47bf491676c..3c3f0c7dcec 100644 --- a/notebooks/v2/stop_instance_test.py +++ b/notebooks/v2/tests/stop_instance_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import stop_instance +from .. import stop_instance PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/update_instance_machine_type_test.py b/notebooks/v2/tests/update_instance_machine_type_test.py similarity index 98% rename from notebooks/v2/update_instance_machine_type_test.py rename to notebooks/v2/tests/update_instance_machine_type_test.py index 6354a8f67fa..885f38f92e5 100644 --- a/notebooks/v2/update_instance_machine_type_test.py +++ b/notebooks/v2/tests/update_instance_machine_type_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import update_instance_machine_type +from .. import update_instance_machine_type PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/update_instance_metadata_test.py b/notebooks/v2/tests/update_instance_metadata_test.py similarity index 98% rename from notebooks/v2/update_instance_metadata_test.py rename to notebooks/v2/tests/update_instance_metadata_test.py index d008f0d01d1..58fcff45393 100644 --- a/notebooks/v2/update_instance_metadata_test.py +++ b/notebooks/v2/tests/update_instance_metadata_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import update_instance_metadata +from .. import update_instance_metadata PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/update_instance_shielded_config_test.py b/notebooks/v2/tests/update_instance_shielded_config_test.py similarity index 98% rename from notebooks/v2/update_instance_shielded_config_test.py rename to notebooks/v2/tests/update_instance_shielded_config_test.py index ddb719b48ef..b3ce03c35f9 100644 --- a/notebooks/v2/update_instance_shielded_config_test.py +++ b/notebooks/v2/tests/update_instance_shielded_config_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import update_instance_shielded_config +from .. import update_instance_shielded_config PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" diff --git a/notebooks/v2/upgrade_instance_test.py b/notebooks/v2/tests/upgrade_instance_test.py similarity index 98% rename from notebooks/v2/upgrade_instance_test.py rename to notebooks/v2/tests/upgrade_instance_test.py index 8232fd8281e..11b4bc8b780 100644 --- a/notebooks/v2/upgrade_instance_test.py +++ b/notebooks/v2/tests/upgrade_instance_test.py @@ -18,7 +18,7 @@ from google.cloud import notebooks_v2 import pytest -import upgrade_instance +from .. import upgrade_instance PROJECT_ID = "example-project" LOCATION_ID = "us-central1-a" From ddbdd065ec938827448f16c53f7e2e02f4cbccdd Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Tue, 3 Nov 2026 19:24:16 +0000 Subject: [PATCH 4/8] test(notebooks): assert on outgoing request payloads in v2 unit tests Problem: The v2 unit tests mostly asserted on stdout, which only echoed fields from the mock response that each test built itself. A sample could drop a security setting (public IP, root access, CMEK, instance owners), send a wrong update mask, hardcode a project or region in a resource path, or skip the IAM read-modify-write append, and its test would still pass. Solution: - Each success-path test now reads the request the sample passed to the mocked client (call_args.kwargs["request"]) and checks resource names, interpolated paths (KMS key, VPC network and subnet, snapshot, service account), security-relevant fields, and update mask paths. - Tests use PROJECT_ID=test-project and LOCATION_ID=us-west1-b instead of the sample defaults, so hardcoded projects or regions fail the test. - set_instance_iam_policy_test checks that the existing binding and etag from get_iam_policy are kept and the new binding is appended, and that set_iam_policy is never called when the read fails. Planted bugs to confirm the new assertions catch them: dropping notebook-disable-root, enabling public IP, skipping the IAM append, hardcoding the region, hardcoding the project in the KMS key and snapshot path, a wrong update mask, and dropping instance_owners. All 8 were caught. --- .../check_instance_iam_permissions_test.py | 8 +++- .../check_instance_upgradability_test.py | 7 +++- .../v2/tests/create_instance_cmek_test.py | 19 ++++++++- .../create_instance_confidential_test.py | 13 +++++- .../create_instance_custom_container_test.py | 13 +++++- .../v2/tests/create_instance_gpu_test.py | 15 ++++++- .../v2/tests/create_instance_hardened_test.py | 41 +++++++++++++++++- .../tests/create_instance_private_vpc_test.py | 18 +++++++- .../tests/create_instance_single_user_test.py | 13 +++++- notebooks/v2/tests/create_instance_test.py | 10 ++++- notebooks/v2/tests/delete_instance_test.py | 7 +++- notebooks/v2/tests/diagnose_instance_test.py | 8 +++- .../v2/tests/get_instance_iam_policy_test.py | 7 +++- notebooks/v2/tests/get_instance_test.py | 7 +++- notebooks/v2/tests/list_instances_test.py | 7 +++- notebooks/v2/tests/reset_instance_test.py | 7 +++- notebooks/v2/tests/rollback_instance_test.py | 11 ++++- .../v2/tests/set_instance_iam_policy_test.py | 42 ++++++++++--------- notebooks/v2/tests/start_instance_test.py | 7 +++- notebooks/v2/tests/stop_instance_test.py | 7 +++- .../update_instance_machine_type_test.py | 9 +++- .../v2/tests/update_instance_metadata_test.py | 10 ++++- .../update_instance_shielded_config_test.py | 18 +++++++- notebooks/v2/tests/upgrade_instance_test.py | 7 +++- 24 files changed, 246 insertions(+), 65 deletions(-) diff --git a/notebooks/v2/tests/check_instance_iam_permissions_test.py b/notebooks/v2/tests/check_instance_iam_permissions_test.py index c15621d946a..5824b887c89 100644 --- a/notebooks/v2/tests/check_instance_iam_permissions_test.py +++ b/notebooks/v2/tests/check_instance_iam_permissions_test.py @@ -20,8 +20,8 @@ from .. import check_instance_iam_permissions -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -42,6 +42,10 @@ def test_check_instance_iam_permissions_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + request = mock_client.test_iam_permissions.call_args.kwargs["request"] + assert request.resource == INSTANCE_NAME + assert list(request.permissions) == ["notebooks.instances.get"] + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Allowed Permissions: ['notebooks.instances.get']" in captured.out diff --git a/notebooks/v2/tests/check_instance_upgradability_test.py b/notebooks/v2/tests/check_instance_upgradability_test.py index 8c81ed1d6cc..053e94464c0 100644 --- a/notebooks/v2/tests/check_instance_upgradability_test.py +++ b/notebooks/v2/tests/check_instance_upgradability_test.py @@ -20,8 +20,8 @@ from .. import check_instance_upgradability -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -45,6 +45,9 @@ def test_check_instance_upgradability_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + request = mock_client.check_instance_upgradability.call_args.kwargs["request"] + assert request.notebook_instance == INSTANCE_NAME + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Upgradeable: True" in captured.out diff --git a/notebooks/v2/tests/create_instance_cmek_test.py b/notebooks/v2/tests/create_instance_cmek_test.py index ccd9cffc52e..1f8e372ab40 100644 --- a/notebooks/v2/tests/create_instance_cmek_test.py +++ b/notebooks/v2/tests/create_instance_cmek_test.py @@ -20,8 +20,8 @@ from .. import create_instance_cmek -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-cmek-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -53,6 +53,21 @@ def test_create_instance_cmek_success( create_instance_cmek.create_instance_cmek(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + region_id = "-".join(LOCATION_ID.split("-")[:-1]) + expected_kms_key = ( + f"projects/{PROJECT_ID}/locations/{region_id}/" + "keyRings/example-key-ring/cryptoKeys/example-key" + ) + cmek = notebooks_v2.DiskEncryption.CMEK + assert gce_setup.boot_disk.disk_encryption == cmek + assert gce_setup.boot_disk.kms_key == expected_kms_key + assert gce_setup.data_disks[0].disk_encryption == cmek + assert gce_setup.data_disks[0].kms_key == expected_kms_key + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Boot Disk Encryption: CMEK" in captured.out diff --git a/notebooks/v2/tests/create_instance_confidential_test.py b/notebooks/v2/tests/create_instance_confidential_test.py index 59fae743b2a..714af1e853b 100644 --- a/notebooks/v2/tests/create_instance_confidential_test.py +++ b/notebooks/v2/tests/create_instance_confidential_test.py @@ -20,8 +20,8 @@ from .. import create_instance_confidential -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-confidential-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -54,6 +54,15 @@ def test_create_instance_confidential_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + shielded = gce_setup.shielded_instance_config + assert shielded.enable_secure_boot + assert shielded.enable_vtpm + assert shielded.enable_integrity_monitoring + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Secure Boot Enabled: True" in captured.out diff --git a/notebooks/v2/tests/create_instance_custom_container_test.py b/notebooks/v2/tests/create_instance_custom_container_test.py index 473d86c794d..55fa094e19a 100644 --- a/notebooks/v2/tests/create_instance_custom_container_test.py +++ b/notebooks/v2/tests/create_instance_custom_container_test.py @@ -20,8 +20,8 @@ from .. import create_instance_custom_container -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-container-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -52,6 +52,15 @@ def test_create_instance_custom_container_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + assert gce_setup.container_image.repository == ( + "gcr.io/deeplearning-platform-release/base-cpu" + ) + assert gce_setup.container_image.tag == "latest" + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert ( diff --git a/notebooks/v2/tests/create_instance_gpu_test.py b/notebooks/v2/tests/create_instance_gpu_test.py index 11e245a0f71..f159d0404b0 100644 --- a/notebooks/v2/tests/create_instance_gpu_test.py +++ b/notebooks/v2/tests/create_instance_gpu_test.py @@ -20,8 +20,8 @@ from .. import create_instance_gpu -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-gpu-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -45,6 +45,17 @@ def test_create_instance_gpu_success( create_instance_gpu.create_instance_gpu(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + accelerator = gce_setup.accelerator_configs[0] + assert accelerator.type_ == ( + notebooks_v2.AcceleratorConfig.AcceleratorType.NVIDIA_TESLA_T4 + ) + assert accelerator.core_count == 1 + assert gce_setup.gpu_driver_config.enable_gpu_driver + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Machine Type: n1-standard-4" in captured.out diff --git a/notebooks/v2/tests/create_instance_hardened_test.py b/notebooks/v2/tests/create_instance_hardened_test.py index 86f0270df89..88c6ebc0f0b 100644 --- a/notebooks/v2/tests/create_instance_hardened_test.py +++ b/notebooks/v2/tests/create_instance_hardened_test.py @@ -20,8 +20,8 @@ from .. import create_instance_hardened -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-hardened-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -57,6 +57,43 @@ def test_create_instance_hardened_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + region_id = "-".join(LOCATION_ID.split("-")[:-1]) + network_interface = gce_setup.network_interfaces[0] + shielded = gce_setup.shielded_instance_config + metadata = dict(gce_setup.metadata) + expected_kms_key = ( + f"projects/{PROJECT_ID}/locations/{region_id}/" + "keyRings/example-key-ring/cryptoKeys/example-key" + ) + assert gce_setup.disable_public_ip + assert network_interface.network == ( + f"projects/{PROJECT_ID}/global/networks/example-network" + ) + assert network_interface.subnet == ( + f"projects/{PROJECT_ID}/regions/{region_id}/subnetworks/example-subnet" + ) + assert gce_setup.service_accounts[0].email == ( + f"example-workbench-sa@{PROJECT_ID}.iam.gserviceaccount.com" + ) + assert list(request.instance.instance_owners) == ["example-user@example.com"] + assert shielded.enable_secure_boot + assert shielded.enable_vtpm + assert shielded.enable_integrity_monitoring + cmek = notebooks_v2.DiskEncryption.CMEK + assert gce_setup.boot_disk.disk_encryption == cmek + assert gce_setup.boot_disk.kms_key == expected_kms_key + assert gce_setup.data_disks[0].disk_encryption == cmek + assert gce_setup.data_disks[0].kms_key == expected_kms_key + assert metadata["notebook-disable-root"] == "true" + assert metadata["notebook-disable-downloads"] == "true" + assert metadata["notebook-disable-terminal"] == "true" + assert metadata["serial-port-enable"] == "false" + assert metadata["block-project-ssh-keys"] == "true" + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Disable Public IP: True" in captured.out diff --git a/notebooks/v2/tests/create_instance_private_vpc_test.py b/notebooks/v2/tests/create_instance_private_vpc_test.py index a31565080f4..2000f72ead7 100644 --- a/notebooks/v2/tests/create_instance_private_vpc_test.py +++ b/notebooks/v2/tests/create_instance_private_vpc_test.py @@ -20,8 +20,8 @@ from .. import create_instance_private_vpc -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-private-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -50,6 +50,20 @@ def test_create_instance_private_vpc_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + region_id = "-".join(LOCATION_ID.split("-")[:-1]) + network_interface = gce_setup.network_interfaces[0] + assert gce_setup.disable_public_ip + assert network_interface.network == ( + f"projects/{PROJECT_ID}/global/networks/example-network" + ) + assert network_interface.subnet == ( + f"projects/{PROJECT_ID}/regions/{region_id}/subnetworks/example-subnet" + ) + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Disable Public IP: True" in captured.out diff --git a/notebooks/v2/tests/create_instance_single_user_test.py b/notebooks/v2/tests/create_instance_single_user_test.py index 8fbbbf38a2b..9d70b9c4eeb 100644 --- a/notebooks/v2/tests/create_instance_single_user_test.py +++ b/notebooks/v2/tests/create_instance_single_user_test.py @@ -20,8 +20,8 @@ from .. import create_instance_single_user -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-single-user-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -48,6 +48,15 @@ def test_create_instance_single_user_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + assert list(request.instance.instance_owners) == ["example-user@example.com"] + assert gce_setup.service_accounts[0].email == ( + f"example-sa@{PROJECT_ID}.iam.gserviceaccount.com" + ) + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Creator: example-user@example.com" in captured.out diff --git a/notebooks/v2/tests/create_instance_test.py b/notebooks/v2/tests/create_instance_test.py index 0ec55a1af3a..0440f68a807 100644 --- a/notebooks/v2/tests/create_instance_test.py +++ b/notebooks/v2/tests/create_instance_test.py @@ -20,8 +20,8 @@ from .. import create_instance -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -46,6 +46,12 @@ def test_create_instance_success( create_instance.create_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.create_instance.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + assert request.instance_id == INSTANCE_ID + gce_setup = request.instance.gce_setup + assert gce_setup.machine_type == "e2-standard-4" + captured = capsys.readouterr() assert "Waiting for operation to complete:" in captured.out assert f"Resource ID: {INSTANCE_NAME}" in captured.out diff --git a/notebooks/v2/tests/delete_instance_test.py b/notebooks/v2/tests/delete_instance_test.py index 9dad8164089..de687a69245 100644 --- a/notebooks/v2/tests/delete_instance_test.py +++ b/notebooks/v2/tests/delete_instance_test.py @@ -19,8 +19,8 @@ from .. import delete_instance -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -41,6 +41,9 @@ def test_delete_instance_success( delete_instance.delete_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.delete_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + captured = capsys.readouterr() assert "Waiting for operation to complete:" in captured.out assert f"Resource ID: {INSTANCE_NAME}" in captured.out diff --git a/notebooks/v2/tests/diagnose_instance_test.py b/notebooks/v2/tests/diagnose_instance_test.py index 5940be65054..31f09fcf04e 100644 --- a/notebooks/v2/tests/diagnose_instance_test.py +++ b/notebooks/v2/tests/diagnose_instance_test.py @@ -20,8 +20,8 @@ from .. import diagnose_instance -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -45,6 +45,10 @@ def test_diagnose_instance_success( diagnose_instance.diagnose_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.diagnose_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + assert request.diagnostic_config.gcs_bucket == "gs://example-bucket" + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Health State: HEALTHY" in captured.out diff --git a/notebooks/v2/tests/get_instance_iam_policy_test.py b/notebooks/v2/tests/get_instance_iam_policy_test.py index 3e9e12538cb..886e5160cd6 100644 --- a/notebooks/v2/tests/get_instance_iam_policy_test.py +++ b/notebooks/v2/tests/get_instance_iam_policy_test.py @@ -20,8 +20,8 @@ from .. import get_instance_iam_policy -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -47,6 +47,9 @@ def test_get_instance_iam_policy_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + request = mock_client.get_iam_policy.call_args.kwargs["request"] + assert request.resource == INSTANCE_NAME + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Role: roles/notebooks.runner" in captured.out diff --git a/notebooks/v2/tests/get_instance_test.py b/notebooks/v2/tests/get_instance_test.py index 1ecd7e6576a..1cdd831662d 100644 --- a/notebooks/v2/tests/get_instance_test.py +++ b/notebooks/v2/tests/get_instance_test.py @@ -20,8 +20,8 @@ from .. import get_instance -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -44,6 +44,9 @@ def test_get_instance_success( get_instance.get_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.get_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Health State: HEALTHY" in captured.out diff --git a/notebooks/v2/tests/list_instances_test.py b/notebooks/v2/tests/list_instances_test.py index 00d7872eafc..8b6c3671977 100644 --- a/notebooks/v2/tests/list_instances_test.py +++ b/notebooks/v2/tests/list_instances_test.py @@ -20,8 +20,8 @@ from .. import list_instances -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -44,6 +44,9 @@ def test_list_instances_success( list_instances.list_instances(PROJECT_ID, LOCATION_ID) + request = mock_client.list_instances.call_args.kwargs["request"] + assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "State: ACTIVE" in captured.out diff --git a/notebooks/v2/tests/reset_instance_test.py b/notebooks/v2/tests/reset_instance_test.py index d0fe3cc455d..bb8589d508b 100644 --- a/notebooks/v2/tests/reset_instance_test.py +++ b/notebooks/v2/tests/reset_instance_test.py @@ -20,8 +20,8 @@ from .. import reset_instance -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -44,6 +44,9 @@ def test_reset_instance_success( reset_instance.reset_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.reset_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + captured = capsys.readouterr() assert "Waiting for operation to complete:" in captured.out assert f"Resource ID: {INSTANCE_NAME}" in captured.out diff --git a/notebooks/v2/tests/rollback_instance_test.py b/notebooks/v2/tests/rollback_instance_test.py index 18215ba8a18..8a01c6cec7c 100644 --- a/notebooks/v2/tests/rollback_instance_test.py +++ b/notebooks/v2/tests/rollback_instance_test.py @@ -20,8 +20,8 @@ from .. import rollback_instance -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -44,6 +44,13 @@ def test_rollback_instance_success( rollback_instance.rollback_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.rollback_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + assert request.target_snapshot == ( + f"projects/{PROJECT_ID}/global/snapshots/example-snapshot" + ) + assert request.revision_id == "1" + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "State: ACTIVE" in captured.out diff --git a/notebooks/v2/tests/set_instance_iam_policy_test.py b/notebooks/v2/tests/set_instance_iam_policy_test.py index bbd35420e36..f6069170520 100644 --- a/notebooks/v2/tests/set_instance_iam_policy_test.py +++ b/notebooks/v2/tests/set_instance_iam_policy_test.py @@ -20,8 +20,8 @@ from .. import set_instance_iam_policy -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -32,33 +32,36 @@ def test_set_instance_iam_policy_success( capsys: pytest.CaptureFixture[str], ) -> None: mock_client = mock_client_cls.return_value - existing_policy = policy_pb2.Policy( - version=1, - etag=b"etag-1", - bindings=[], + existing_binding = policy_pb2.Binding( + role="roles/notebooks.admin", + members=["user:existing-admin@example.com"], ) - updated_policy = policy_pb2.Policy( + mock_client.get_iam_policy.return_value = policy_pb2.Policy( version=1, - etag=b"etag-2", - bindings=[ - policy_pb2.Binding( - role="roles/notebooks.runner", - members=["user:example-user@example.com"], - ) - ], + etag=b"etag-1", + bindings=[existing_binding], ) - mock_client.get_iam_policy.return_value = existing_policy - mock_client.set_iam_policy.return_value = updated_policy + mock_client.set_iam_policy.side_effect = lambda request: request.policy set_instance_iam_policy.set_instance_iam_policy( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) - mock_client.get_iam_policy.assert_called_once() - mock_client.set_iam_policy.assert_called_once() + get_request = mock_client.get_iam_policy.call_args.kwargs["request"] + assert get_request.resource == INSTANCE_NAME + set_request = mock_client.set_iam_policy.call_args.kwargs["request"] + assert set_request.resource == INSTANCE_NAME + # The etag from the read must be sent back so concurrent edits are rejected. + assert set_request.policy.etag == b"etag-1" + bindings = {(b.role, tuple(b.members)) for b in set_request.policy.bindings} + assert bindings == { + ("roles/notebooks.admin", ("user:existing-admin@example.com",)), + ("roles/notebooks.runner", ("user:example-user@example.com",)), + } + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out - assert "Bindings Count: 1" in captured.out + assert "Bindings Count: 2" in captured.out assert captured.err == "" @@ -74,6 +77,7 @@ def test_set_instance_iam_policy_not_found( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + mock_client.set_iam_policy.assert_not_called() captured = capsys.readouterr() assert "error:" in captured.err assert INSTANCE_ID in captured.err diff --git a/notebooks/v2/tests/start_instance_test.py b/notebooks/v2/tests/start_instance_test.py index 9ad9158422f..0c3509af5ca 100644 --- a/notebooks/v2/tests/start_instance_test.py +++ b/notebooks/v2/tests/start_instance_test.py @@ -20,8 +20,8 @@ from .. import start_instance -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -45,6 +45,9 @@ def test_start_instance_success( start_instance.start_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.start_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + captured = capsys.readouterr() assert "Waiting for operation to complete:" in captured.out assert f"Resource ID: {INSTANCE_NAME}" in captured.out diff --git a/notebooks/v2/tests/stop_instance_test.py b/notebooks/v2/tests/stop_instance_test.py index 3c3f0c7dcec..63a69682a79 100644 --- a/notebooks/v2/tests/stop_instance_test.py +++ b/notebooks/v2/tests/stop_instance_test.py @@ -20,8 +20,8 @@ from .. import stop_instance -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -44,6 +44,9 @@ def test_stop_instance_success( stop_instance.stop_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.stop_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + captured = capsys.readouterr() assert "Waiting for operation to complete:" in captured.out assert f"Resource ID: {INSTANCE_NAME}" in captured.out diff --git a/notebooks/v2/tests/update_instance_machine_type_test.py b/notebooks/v2/tests/update_instance_machine_type_test.py index 885f38f92e5..c282c449198 100644 --- a/notebooks/v2/tests/update_instance_machine_type_test.py +++ b/notebooks/v2/tests/update_instance_machine_type_test.py @@ -20,8 +20,8 @@ from .. import update_instance_machine_type -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -47,6 +47,11 @@ def test_update_instance_machine_type_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID, "n1-standard-8" ) + request = mock_client.update_instance.call_args.kwargs["request"] + assert request.instance.name == INSTANCE_NAME + assert request.instance.gce_setup.machine_type == "n1-standard-8" + assert list(request.update_mask.paths) == ["gce_setup.machine_type"] + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Machine Type: n1-standard-8" in captured.out diff --git a/notebooks/v2/tests/update_instance_metadata_test.py b/notebooks/v2/tests/update_instance_metadata_test.py index 58fcff45393..073f20b38fe 100644 --- a/notebooks/v2/tests/update_instance_metadata_test.py +++ b/notebooks/v2/tests/update_instance_metadata_test.py @@ -20,8 +20,8 @@ from .. import update_instance_metadata -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -47,6 +47,12 @@ def test_update_instance_metadata_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + request = mock_client.update_instance.call_args.kwargs["request"] + assert request.instance.name == INSTANCE_NAME + metadata = dict(request.instance.gce_setup.metadata) + assert metadata == {"idle-timeout-seconds": "3600"} + assert list(request.update_mask.paths) == ["gce_setup.metadata"] + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "idle-timeout-seconds" in captured.out diff --git a/notebooks/v2/tests/update_instance_shielded_config_test.py b/notebooks/v2/tests/update_instance_shielded_config_test.py index b3ce03c35f9..413a2a75801 100644 --- a/notebooks/v2/tests/update_instance_shielded_config_test.py +++ b/notebooks/v2/tests/update_instance_shielded_config_test.py @@ -20,8 +20,8 @@ from .. import update_instance_shielded_config -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -53,6 +53,20 @@ def test_update_instance_shielded_config_success( PROJECT_ID, LOCATION_ID, INSTANCE_ID ) + request = mock_client.update_instance.call_args.kwargs["request"] + assert request.instance.name == INSTANCE_NAME + shielded = request.instance.gce_setup.shielded_instance_config + assert shielded.enable_secure_boot + assert shielded.enable_vtpm + assert shielded.enable_integrity_monitoring + assert sorted(request.update_mask.paths) == sorted( + [ + "gce_setup.shielded_instance_config.enable_secure_boot", + "gce_setup.shielded_instance_config.enable_vtpm", + "gce_setup.shielded_instance_config.enable_integrity_monitoring", + ] + ) + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "Secure Boot Enabled: True" in captured.out diff --git a/notebooks/v2/tests/upgrade_instance_test.py b/notebooks/v2/tests/upgrade_instance_test.py index 11b4bc8b780..130d8c08fea 100644 --- a/notebooks/v2/tests/upgrade_instance_test.py +++ b/notebooks/v2/tests/upgrade_instance_test.py @@ -20,8 +20,8 @@ from .. import upgrade_instance -PROJECT_ID = "example-project" -LOCATION_ID = "us-central1-a" +PROJECT_ID = "test-project" +LOCATION_ID = "us-west1-b" INSTANCE_ID = "example-instance" INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" @@ -44,6 +44,9 @@ def test_upgrade_instance_success( upgrade_instance.upgrade_instance(PROJECT_ID, LOCATION_ID, INSTANCE_ID) + request = mock_client.upgrade_instance.call_args.kwargs["request"] + assert request.name == INSTANCE_NAME + captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out assert "State: ACTIVE" in captured.out From 65e7af6b5b61c7aead3353bc830a633db9fc560b Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Tue, 3 Nov 2026 20:46:27 +0000 Subject: [PATCH 5/8] fix(notebooks): clarify LRO completion and tidy v2 sample output - Create and start samples: note that the operation completes before the instance is ready, and point to get_instance to wait for ACTIVE. - update_instance_metadata: print the updated key, the image version, and the metadata key count instead of the full metadata map. - IAM policy samples: print the etag as base64 instead of raw bytes. --- notebooks/v2/create_instance.py | 6 ++++++ notebooks/v2/create_instance_cmek.py | 6 ++++++ notebooks/v2/create_instance_confidential.py | 6 ++++++ notebooks/v2/create_instance_custom_container.py | 6 ++++++ notebooks/v2/create_instance_gpu.py | 6 ++++++ notebooks/v2/create_instance_hardened.py | 6 ++++++ notebooks/v2/create_instance_private_vpc.py | 6 ++++++ notebooks/v2/create_instance_single_user.py | 6 ++++++ notebooks/v2/get_instance_iam_policy.py | 5 ++++- notebooks/v2/set_instance_iam_policy.py | 5 ++++- notebooks/v2/start_instance.py | 6 ++++++ notebooks/v2/tests/get_instance_iam_policy_test.py | 1 + notebooks/v2/tests/set_instance_iam_policy_test.py | 1 + notebooks/v2/tests/update_instance_metadata_test.py | 13 ++++++++++--- notebooks/v2/update_instance_metadata.py | 8 ++++++-- 15 files changed, 80 insertions(+), 7 deletions(-) diff --git a/notebooks/v2/create_instance.py b/notebooks/v2/create_instance.py index 100edd83fb9..b3e3e65a3f7 100644 --- a/notebooks/v2/create_instance.py +++ b/notebooks/v2/create_instance.py @@ -44,12 +44,18 @@ def create_instance( try: operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. response = operation.result() print(f"Resource ID: {response.name}") print(f"Creation Date: {response.create_time}") print(f"Machine Type: {response.gce_setup.machine_type}") print(f"Proxy URI: {response.proxy_uri}") print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) except exceptions.AlreadyExists as err: print( f"error: Instance '{instance_id}' already exists in location " diff --git a/notebooks/v2/create_instance_cmek.py b/notebooks/v2/create_instance_cmek.py index 3891cbf07bb..0203ae7c2f3 100644 --- a/notebooks/v2/create_instance_cmek.py +++ b/notebooks/v2/create_instance_cmek.py @@ -70,12 +70,18 @@ def create_instance_cmek( try: operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. response = operation.result() boot_encryption = response.gce_setup.boot_disk.disk_encryption.name print(f"Resource ID: {response.name}") print(f"Creation Date: {response.create_time}") print(f"Boot Disk Encryption: {boot_encryption}") print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) except exceptions.AlreadyExists as err: print( f"error: Instance '{instance_id}' already exists in location " diff --git a/notebooks/v2/create_instance_confidential.py b/notebooks/v2/create_instance_confidential.py index d88d79f3286..ca35fb819a0 100644 --- a/notebooks/v2/create_instance_confidential.py +++ b/notebooks/v2/create_instance_confidential.py @@ -52,12 +52,18 @@ def create_instance_confidential( try: operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. response = operation.result() secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot print(f"Resource ID: {response.name}") print(f"Machine Type: {response.gce_setup.machine_type}") print(f"Secure Boot Enabled: {secure_boot}") print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) except exceptions.AlreadyExists as err: print( f"error: Instance '{instance_id}' already exists in location " diff --git a/notebooks/v2/create_instance_custom_container.py b/notebooks/v2/create_instance_custom_container.py index bec6cf469d7..038d9e33e0c 100644 --- a/notebooks/v2/create_instance_custom_container.py +++ b/notebooks/v2/create_instance_custom_container.py @@ -53,11 +53,17 @@ def create_instance_custom_container( try: operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. response = operation.result() repo = response.gce_setup.container_image.repository print(f"Resource ID: {response.name}") print(f"Container Repository: {repo}") print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) except exceptions.AlreadyExists as err: print( f"error: Instance '{instance_id}' already exists in location " diff --git a/notebooks/v2/create_instance_gpu.py b/notebooks/v2/create_instance_gpu.py index 76aad0cd1f7..6925bd75dc5 100644 --- a/notebooks/v2/create_instance_gpu.py +++ b/notebooks/v2/create_instance_gpu.py @@ -55,11 +55,17 @@ def create_instance_gpu( try: operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. response = operation.result() print(f"Resource ID: {response.name}") print(f"Creation Date: {response.create_time}") print(f"Machine Type: {response.gce_setup.machine_type}") print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) except exceptions.AlreadyExists as err: print( f"error: Instance '{instance_id}' already exists in location " diff --git a/notebooks/v2/create_instance_hardened.py b/notebooks/v2/create_instance_hardened.py index dc355c79214..000e9d788c3 100644 --- a/notebooks/v2/create_instance_hardened.py +++ b/notebooks/v2/create_instance_hardened.py @@ -115,6 +115,8 @@ def create_instance_hardened( try: operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. response = operation.result() secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot print(f"Resource ID: {response.name}") @@ -122,6 +124,10 @@ def create_instance_hardened( print(f"Disable Public IP: {response.gce_setup.disable_public_ip}") print(f"Secure Boot Enabled: {secure_boot}") print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) except exceptions.AlreadyExists as err: print( f"error: Instance '{instance_id}' already exists in location " diff --git a/notebooks/v2/create_instance_private_vpc.py b/notebooks/v2/create_instance_private_vpc.py index c8f7c8604b4..d077aff22fb 100644 --- a/notebooks/v2/create_instance_private_vpc.py +++ b/notebooks/v2/create_instance_private_vpc.py @@ -65,12 +65,18 @@ def create_instance_private_vpc( try: operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. response = operation.result() print(f"Resource ID: {response.name}") print(f"Creation Date: {response.create_time}") print(f"Disable Public IP: {response.gce_setup.disable_public_ip}") print(f"Machine Type: {response.gce_setup.machine_type}") print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) except exceptions.AlreadyExists as err: print( f"error: Instance '{instance_id}' already exists in location " diff --git a/notebooks/v2/create_instance_single_user.py b/notebooks/v2/create_instance_single_user.py index f5077073525..eb838af81b3 100644 --- a/notebooks/v2/create_instance_single_user.py +++ b/notebooks/v2/create_instance_single_user.py @@ -56,10 +56,16 @@ def create_instance_single_user( try: operation = client.create_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes once the instance is accepted, before the + # VM finishes provisioning, so the state is usually PROVISIONING here. response = operation.result() print(f"Resource ID: {response.name}") print(f"Creator: {response.creator}") print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) except exceptions.AlreadyExists as err: print( f"error: Instance '{instance_id}' already exists in location " diff --git a/notebooks/v2/get_instance_iam_policy.py b/notebooks/v2/get_instance_iam_policy.py index 8c8d3660759..b42559951e8 100644 --- a/notebooks/v2/get_instance_iam_policy.py +++ b/notebooks/v2/get_instance_iam_policy.py @@ -13,6 +13,7 @@ # limitations under the License. # [START notebooks_v2_get_instance_iam_policy] +import base64 import sys from google.api_core import exceptions @@ -39,7 +40,9 @@ def get_instance_iam_policy( try: response = client.get_iam_policy(request=request) print(f"Resource ID: {name}") - print(f"Policy Etag: {response.etag}") + # The etag is an opaque byte string; base64 matches how gcloud and + # the REST API display it. + print(f"Policy Etag: {base64.b64encode(response.etag).decode()}") print(f"Policy Version: {response.version}") for binding in response.bindings: print(f"Role: {binding.role}") diff --git a/notebooks/v2/set_instance_iam_policy.py b/notebooks/v2/set_instance_iam_policy.py index 998278bfb4f..3ad73155b2f 100644 --- a/notebooks/v2/set_instance_iam_policy.py +++ b/notebooks/v2/set_instance_iam_policy.py @@ -13,6 +13,7 @@ # limitations under the License. # [START notebooks_v2_set_instance_iam_policy] +import base64 import sys from google.api_core import exceptions @@ -59,7 +60,9 @@ def set_instance_iam_policy( response = client.set_iam_policy(request=set_request) print(f"Resource ID: {name}") print(f"Bindings Count: {len(response.bindings)}") - print(f"Policy Etag: {response.etag}") + # The etag is an opaque byte string; base64 matches how gcloud and + # the REST API display it. + print(f"Policy Etag: {base64.b64encode(response.etag).decode()}") print(f"Policy Version: {response.version}") except exceptions.NotFound as err: print( diff --git a/notebooks/v2/start_instance.py b/notebooks/v2/start_instance.py index 0a4ef2f64d3..880d5813a5b 100644 --- a/notebooks/v2/start_instance.py +++ b/notebooks/v2/start_instance.py @@ -38,10 +38,16 @@ def start_instance( try: operation = client.start_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") + # The operation completes before the VM finishes booting, so the + # state is usually PROVISIONING here. response = operation.result() print(f"Resource ID: {response.name}") print(f"Proxy URI: {response.proxy_uri}") print(f"State: {response.state.name}") + print( + "The instance is still starting up; call get_instance until the " + "state is ACTIVE before connecting to JupyterLab." + ) except exceptions.NotFound as err: print( f"error: Instance '{instance_id}' not found in location " diff --git a/notebooks/v2/tests/get_instance_iam_policy_test.py b/notebooks/v2/tests/get_instance_iam_policy_test.py index 886e5160cd6..9ac7203c2f0 100644 --- a/notebooks/v2/tests/get_instance_iam_policy_test.py +++ b/notebooks/v2/tests/get_instance_iam_policy_test.py @@ -52,6 +52,7 @@ def test_get_instance_iam_policy_success( captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Policy Etag: ZXRhZy0x" in captured.out assert "Role: roles/notebooks.runner" in captured.out assert captured.err == "" diff --git a/notebooks/v2/tests/set_instance_iam_policy_test.py b/notebooks/v2/tests/set_instance_iam_policy_test.py index f6069170520..57df810a9a4 100644 --- a/notebooks/v2/tests/set_instance_iam_policy_test.py +++ b/notebooks/v2/tests/set_instance_iam_policy_test.py @@ -61,6 +61,7 @@ def test_set_instance_iam_policy_success( captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out + assert "Policy Etag: ZXRhZy0x" in captured.out assert "Bindings Count: 2" in captured.out assert captured.err == "" diff --git a/notebooks/v2/tests/update_instance_metadata_test.py b/notebooks/v2/tests/update_instance_metadata_test.py index 073f20b38fe..9a20decb705 100644 --- a/notebooks/v2/tests/update_instance_metadata_test.py +++ b/notebooks/v2/tests/update_instance_metadata_test.py @@ -39,7 +39,13 @@ def test_update_instance_metadata_success( operation.result.return_value = notebooks_v2.Instance( name=INSTANCE_NAME, state=notebooks_v2.State.ACTIVE, - gce_setup=notebooks_v2.GceSetup(metadata={"idle-timeout-seconds": "3600"}), + gce_setup=notebooks_v2.GceSetup( + metadata={ + "idle-timeout-seconds": "3600", + "version": "20261101-2130", + "notebooks-api-version": "v2", + } + ), ) mock_client.update_instance.return_value = operation @@ -55,8 +61,9 @@ def test_update_instance_metadata_success( captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out - assert "idle-timeout-seconds" in captured.out - assert "State: ACTIVE" in captured.out + assert "Idle Timeout Seconds: 3600" in captured.out + assert "Image Version: 20261101-2130" in captured.out + assert "Metadata Key Count: 3" in captured.out assert captured.err == "" diff --git a/notebooks/v2/update_instance_metadata.py b/notebooks/v2/update_instance_metadata.py index 36067ff3066..e6e1eac804b 100644 --- a/notebooks/v2/update_instance_metadata.py +++ b/notebooks/v2/update_instance_metadata.py @@ -48,9 +48,13 @@ def update_instance_metadata( operation = client.update_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") response = operation.result() + # The service merges the update into metadata it manages itself, such + # as the image version, so only the updated key changes. + metadata = response.gce_setup.metadata print(f"Resource ID: {response.name}") - print(f"Metadata: {response.gce_setup.metadata}") - print(f"State: {response.state.name}") + print(f"Idle Timeout Seconds: {metadata.get('idle-timeout-seconds')}") + print(f"Image Version: {metadata.get('version')}") + print(f"Metadata Key Count: {len(metadata)}") except exceptions.NotFound as err: print( f"error: Instance '{instance_id}' not found in location " From c9041d3290bd33c3bd4d49365777b7d73ca8fa38 Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Thu, 5 Nov 2026 20:23:47 +0000 Subject: [PATCH 6/8] fix(notebooks): accept full service_account_email in single-user and hardened samples Constructing service_account_email dynamically from project_id prevented callers from passing a Default Compute Engine service account or a cross-project service account. Updates create_instance_single_user.py and create_instance_hardened.py (and their unit tests) to accept the full service_account_email directly. --- notebooks/v2/create_instance_hardened.py | 6 +++--- notebooks/v2/create_instance_single_user.py | 6 +++--- notebooks/v2/tests/create_instance_hardened_test.py | 2 +- notebooks/v2/tests/create_instance_single_user_test.py | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/notebooks/v2/create_instance_hardened.py b/notebooks/v2/create_instance_hardened.py index 000e9d788c3..0cd6774694d 100644 --- a/notebooks/v2/create_instance_hardened.py +++ b/notebooks/v2/create_instance_hardened.py @@ -25,7 +25,7 @@ def create_instance_hardened( instance_id: str = "example-hardened-instance", network_id: str = "example-network", subnet_id: str = "example-subnet", - service_account_id: str = "example-workbench-sa", + service_account_email: str = "example-workbench-sa@example-project-id.iam.gserviceaccount.com", instance_owner: str = "example-user@example.com", key_ring_id: str = "example-key-ring", crypto_key_id: str = "example-key", @@ -44,7 +44,8 @@ def create_instance_hardened( instance_id: The ID to assign to the new Workbench instance. network_id: The ID of the VPC network in the project. subnet_id: The ID of the subnetwork in the instance region. - service_account_id: The ID of the custom service account in the project. + service_account_email: The email address of the service account to + attach to the instance. instance_owner: The single user email permitted to access the instance. key_ring_id: The Cloud KMS key ring ID in the instance region. crypto_key_id: The Cloud KMS crypto key ID used for disk encryption. @@ -55,7 +56,6 @@ def create_instance_hardened( region_id = "-".join(parts[:-1]) if len(parts) > 2 else location_id network = f"projects/{project_id}/global/networks/{network_id}" subnet = f"projects/{project_id}/regions/{region_id}/subnetworks/{subnet_id}" - service_account_email = f"{service_account_id}@{project_id}.iam.gserviceaccount.com" kms_key = ( f"projects/{project_id}/locations/{region_id}/" f"keyRings/{key_ring_id}/cryptoKeys/{crypto_key_id}" diff --git a/notebooks/v2/create_instance_single_user.py b/notebooks/v2/create_instance_single_user.py index eb838af81b3..73b30402a9a 100644 --- a/notebooks/v2/create_instance_single_user.py +++ b/notebooks/v2/create_instance_single_user.py @@ -24,7 +24,7 @@ def create_instance_single_user( location_id: str = "us-central1-a", instance_id: str = "example-single-user-instance", instance_owner: str = "example-user@example.com", - service_account_id: str = "example-sa", + service_account_email: str = "example-sa@example-project-id.iam.gserviceaccount.com", ) -> None: """Creates a single-user Workbench instance with a custom service account. @@ -33,11 +33,11 @@ def create_instance_single_user( location_id: The Compute Engine zone for the new instance. instance_id: The ID to assign to the new Workbench instance. instance_owner: The email address of the single permitted user. - service_account_id: The ID of the custom service account in the project. + service_account_email: The email address of the service account to + attach to the instance. """ client = notebooks_v2.NotebookServiceClient() parent = f"projects/{project_id}/locations/{location_id}" - service_account_email = f"{service_account_id}@{project_id}.iam.gserviceaccount.com" service_account = notebooks_v2.ServiceAccount(email=service_account_email) gce_setup = notebooks_v2.GceSetup( machine_type="e2-standard-4", diff --git a/notebooks/v2/tests/create_instance_hardened_test.py b/notebooks/v2/tests/create_instance_hardened_test.py index 88c6ebc0f0b..9048b973bdf 100644 --- a/notebooks/v2/tests/create_instance_hardened_test.py +++ b/notebooks/v2/tests/create_instance_hardened_test.py @@ -77,7 +77,7 @@ def test_create_instance_hardened_success( f"projects/{PROJECT_ID}/regions/{region_id}/subnetworks/example-subnet" ) assert gce_setup.service_accounts[0].email == ( - f"example-workbench-sa@{PROJECT_ID}.iam.gserviceaccount.com" + "example-workbench-sa@example-project-id.iam.gserviceaccount.com" ) assert list(request.instance.instance_owners) == ["example-user@example.com"] assert shielded.enable_secure_boot diff --git a/notebooks/v2/tests/create_instance_single_user_test.py b/notebooks/v2/tests/create_instance_single_user_test.py index 9d70b9c4eeb..f8b2c53dce3 100644 --- a/notebooks/v2/tests/create_instance_single_user_test.py +++ b/notebooks/v2/tests/create_instance_single_user_test.py @@ -54,7 +54,7 @@ def test_create_instance_single_user_success( gce_setup = request.instance.gce_setup assert list(request.instance.instance_owners) == ["example-user@example.com"] assert gce_setup.service_accounts[0].email == ( - f"example-sa@{PROJECT_ID}.iam.gserviceaccount.com" + "example-sa@example-project-id.iam.gserviceaccount.com" ) captured = capsys.readouterr() From eccd1c6b841075bae9a96d1cce80ddd43f6246af Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Thu, 5 Nov 2026 22:36:41 +0000 Subject: [PATCH 7/8] test(notebooks): use target-safe mock.patch.object with autospec in v2 unit tests Replace raw string paths in @mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") across all 24 unit test files in notebooks/v2/tests/ with target-safe @mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True). Also pass explicit service_account_email arguments in create_instance_single_user_test.py and create_instance_hardened_test.py to verify non-default service account email forwarding. --- .../tests/check_instance_iam_permissions_test.py | 5 +++-- .../v2/tests/check_instance_upgradability_test.py | 4 ++-- notebooks/v2/tests/create_instance_cmek_test.py | 4 ++-- .../v2/tests/create_instance_confidential_test.py | 4 ++-- .../tests/create_instance_custom_container_test.py | 4 ++-- notebooks/v2/tests/create_instance_gpu_test.py | 4 ++-- .../v2/tests/create_instance_hardened_test.py | 14 ++++++++------ .../v2/tests/create_instance_private_vpc_test.py | 4 ++-- .../v2/tests/create_instance_single_user_test.py | 14 ++++++++------ notebooks/v2/tests/create_instance_test.py | 4 ++-- notebooks/v2/tests/delete_instance_test.py | 5 +++-- notebooks/v2/tests/diagnose_instance_test.py | 4 ++-- notebooks/v2/tests/get_instance_iam_policy_test.py | 5 +++-- notebooks/v2/tests/get_instance_test.py | 4 ++-- notebooks/v2/tests/list_instances_test.py | 4 ++-- notebooks/v2/tests/reset_instance_test.py | 4 ++-- notebooks/v2/tests/rollback_instance_test.py | 4 ++-- notebooks/v2/tests/set_instance_iam_policy_test.py | 5 +++-- notebooks/v2/tests/start_instance_test.py | 4 ++-- notebooks/v2/tests/stop_instance_test.py | 4 ++-- .../v2/tests/update_instance_machine_type_test.py | 4 ++-- .../v2/tests/update_instance_metadata_test.py | 4 ++-- .../tests/update_instance_shielded_config_test.py | 4 ++-- notebooks/v2/tests/upgrade_instance_test.py | 4 ++-- 24 files changed, 64 insertions(+), 56 deletions(-) diff --git a/notebooks/v2/tests/check_instance_iam_permissions_test.py b/notebooks/v2/tests/check_instance_iam_permissions_test.py index 5824b887c89..797128f409e 100644 --- a/notebooks/v2/tests/check_instance_iam_permissions_test.py +++ b/notebooks/v2/tests/check_instance_iam_permissions_test.py @@ -15,6 +15,7 @@ from unittest import mock from google.api_core import exceptions +from google.cloud import notebooks_v2 from google.iam.v1 import iam_policy_pb2 import pytest @@ -26,7 +27,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_check_instance_iam_permissions_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -52,7 +53,7 @@ def test_check_instance_iam_permissions_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_check_instance_iam_permissions_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/check_instance_upgradability_test.py b/notebooks/v2/tests/check_instance_upgradability_test.py index 053e94464c0..1eb4919c3dd 100644 --- a/notebooks/v2/tests/check_instance_upgradability_test.py +++ b/notebooks/v2/tests/check_instance_upgradability_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_check_instance_upgradability_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -54,7 +54,7 @@ def test_check_instance_upgradability_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_check_instance_upgradability_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/create_instance_cmek_test.py b/notebooks/v2/tests/create_instance_cmek_test.py index 1f8e372ab40..c1188c5acec 100644 --- a/notebooks/v2/tests/create_instance_cmek_test.py +++ b/notebooks/v2/tests/create_instance_cmek_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_cmek_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -75,7 +75,7 @@ def test_create_instance_cmek_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_cmek_already_exists( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/create_instance_confidential_test.py b/notebooks/v2/tests/create_instance_confidential_test.py index 714af1e853b..a3b7c8a32b6 100644 --- a/notebooks/v2/tests/create_instance_confidential_test.py +++ b/notebooks/v2/tests/create_instance_confidential_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_confidential_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -70,7 +70,7 @@ def test_create_instance_confidential_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_confidential_already_exists( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/create_instance_custom_container_test.py b/notebooks/v2/tests/create_instance_custom_container_test.py index 55fa094e19a..9a172d133c0 100644 --- a/notebooks/v2/tests/create_instance_custom_container_test.py +++ b/notebooks/v2/tests/create_instance_custom_container_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_custom_container_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -71,7 +71,7 @@ def test_create_instance_custom_container_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_custom_container_already_exists( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/create_instance_gpu_test.py b/notebooks/v2/tests/create_instance_gpu_test.py index f159d0404b0..3da54cb01df 100644 --- a/notebooks/v2/tests/create_instance_gpu_test.py +++ b/notebooks/v2/tests/create_instance_gpu_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_gpu_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -63,7 +63,7 @@ def test_create_instance_gpu_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_gpu_already_exists( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/create_instance_hardened_test.py b/notebooks/v2/tests/create_instance_hardened_test.py index 9048b973bdf..de847e7fe49 100644 --- a/notebooks/v2/tests/create_instance_hardened_test.py +++ b/notebooks/v2/tests/create_instance_hardened_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_hardened_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -53,8 +53,12 @@ def test_create_instance_hardened_success( ) mock_client.create_instance.return_value = operation + custom_sa_email = "custom-workbench-sa@shared-project.iam.gserviceaccount.com" create_instance_hardened.create_instance_hardened( - PROJECT_ID, LOCATION_ID, INSTANCE_ID + PROJECT_ID, + LOCATION_ID, + INSTANCE_ID, + service_account_email=custom_sa_email, ) request = mock_client.create_instance.call_args.kwargs["request"] @@ -76,9 +80,7 @@ def test_create_instance_hardened_success( assert network_interface.subnet == ( f"projects/{PROJECT_ID}/regions/{region_id}/subnetworks/example-subnet" ) - assert gce_setup.service_accounts[0].email == ( - "example-workbench-sa@example-project-id.iam.gserviceaccount.com" - ) + assert gce_setup.service_accounts[0].email == custom_sa_email assert list(request.instance.instance_owners) == ["example-user@example.com"] assert shielded.enable_secure_boot assert shielded.enable_vtpm @@ -102,7 +104,7 @@ def test_create_instance_hardened_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_hardened_already_exists( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/create_instance_private_vpc_test.py b/notebooks/v2/tests/create_instance_private_vpc_test.py index 2000f72ead7..325275a7867 100644 --- a/notebooks/v2/tests/create_instance_private_vpc_test.py +++ b/notebooks/v2/tests/create_instance_private_vpc_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_private_vpc_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -71,7 +71,7 @@ def test_create_instance_private_vpc_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_private_vpc_already_exists( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/create_instance_single_user_test.py b/notebooks/v2/tests/create_instance_single_user_test.py index f8b2c53dce3..521c503691f 100644 --- a/notebooks/v2/tests/create_instance_single_user_test.py +++ b/notebooks/v2/tests/create_instance_single_user_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_single_user_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -44,8 +44,12 @@ def test_create_instance_single_user_success( ) mock_client.create_instance.return_value = operation + custom_sa_email = "123456789012-compute@developer.gserviceaccount.com" create_instance_single_user.create_instance_single_user( - PROJECT_ID, LOCATION_ID, INSTANCE_ID + PROJECT_ID, + LOCATION_ID, + INSTANCE_ID, + service_account_email=custom_sa_email, ) request = mock_client.create_instance.call_args.kwargs["request"] @@ -53,9 +57,7 @@ def test_create_instance_single_user_success( assert request.instance_id == INSTANCE_ID gce_setup = request.instance.gce_setup assert list(request.instance.instance_owners) == ["example-user@example.com"] - assert gce_setup.service_accounts[0].email == ( - "example-sa@example-project-id.iam.gserviceaccount.com" - ) + assert gce_setup.service_accounts[0].email == custom_sa_email captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out @@ -64,7 +66,7 @@ def test_create_instance_single_user_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_single_user_already_exists( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/create_instance_test.py b/notebooks/v2/tests/create_instance_test.py index 0440f68a807..e9d77c2f7ee 100644 --- a/notebooks/v2/tests/create_instance_test.py +++ b/notebooks/v2/tests/create_instance_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -60,7 +60,7 @@ def test_create_instance_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_create_instance_already_exists( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/delete_instance_test.py b/notebooks/v2/tests/delete_instance_test.py index de687a69245..66192e9e9d8 100644 --- a/notebooks/v2/tests/delete_instance_test.py +++ b/notebooks/v2/tests/delete_instance_test.py @@ -15,6 +15,7 @@ from unittest import mock from google.api_core import exceptions +from google.cloud import notebooks_v2 import pytest from .. import delete_instance @@ -25,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_delete_instance_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -51,7 +52,7 @@ def test_delete_instance_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_delete_instance_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/diagnose_instance_test.py b/notebooks/v2/tests/diagnose_instance_test.py index 31f09fcf04e..6388cc1cf95 100644 --- a/notebooks/v2/tests/diagnose_instance_test.py +++ b/notebooks/v2/tests/diagnose_instance_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_diagnose_instance_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -56,7 +56,7 @@ def test_diagnose_instance_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_diagnose_instance_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/get_instance_iam_policy_test.py b/notebooks/v2/tests/get_instance_iam_policy_test.py index 9ac7203c2f0..edbc0e5ca69 100644 --- a/notebooks/v2/tests/get_instance_iam_policy_test.py +++ b/notebooks/v2/tests/get_instance_iam_policy_test.py @@ -15,6 +15,7 @@ from unittest import mock from google.api_core import exceptions +from google.cloud import notebooks_v2 from google.iam.v1 import policy_pb2 import pytest @@ -26,7 +27,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_get_instance_iam_policy_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -57,7 +58,7 @@ def test_get_instance_iam_policy_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_get_instance_iam_policy_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/get_instance_test.py b/notebooks/v2/tests/get_instance_test.py index 1cdd831662d..74f86ea9b29 100644 --- a/notebooks/v2/tests/get_instance_test.py +++ b/notebooks/v2/tests/get_instance_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_get_instance_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -54,7 +54,7 @@ def test_get_instance_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_get_instance_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/list_instances_test.py b/notebooks/v2/tests/list_instances_test.py index 8b6c3671977..d651ca58dc0 100644 --- a/notebooks/v2/tests/list_instances_test.py +++ b/notebooks/v2/tests/list_instances_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_list_instances_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -53,7 +53,7 @@ def test_list_instances_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_list_instances_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/reset_instance_test.py b/notebooks/v2/tests/reset_instance_test.py index bb8589d508b..3b99a450325 100644 --- a/notebooks/v2/tests/reset_instance_test.py +++ b/notebooks/v2/tests/reset_instance_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_reset_instance_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -54,7 +54,7 @@ def test_reset_instance_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_reset_instance_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/rollback_instance_test.py b/notebooks/v2/tests/rollback_instance_test.py index 8a01c6cec7c..ded7cd16a40 100644 --- a/notebooks/v2/tests/rollback_instance_test.py +++ b/notebooks/v2/tests/rollback_instance_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_rollback_instance_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -57,7 +57,7 @@ def test_rollback_instance_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_rollback_instance_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/set_instance_iam_policy_test.py b/notebooks/v2/tests/set_instance_iam_policy_test.py index 57df810a9a4..bf1a2e5ebbc 100644 --- a/notebooks/v2/tests/set_instance_iam_policy_test.py +++ b/notebooks/v2/tests/set_instance_iam_policy_test.py @@ -15,6 +15,7 @@ from unittest import mock from google.api_core import exceptions +from google.cloud import notebooks_v2 from google.iam.v1 import policy_pb2 import pytest @@ -26,7 +27,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_set_instance_iam_policy_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -66,7 +67,7 @@ def test_set_instance_iam_policy_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_set_instance_iam_policy_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/start_instance_test.py b/notebooks/v2/tests/start_instance_test.py index 0c3509af5ca..0f8e0cb10a2 100644 --- a/notebooks/v2/tests/start_instance_test.py +++ b/notebooks/v2/tests/start_instance_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_start_instance_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -55,7 +55,7 @@ def test_start_instance_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_start_instance_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/stop_instance_test.py b/notebooks/v2/tests/stop_instance_test.py index 63a69682a79..8b66c0dd0b8 100644 --- a/notebooks/v2/tests/stop_instance_test.py +++ b/notebooks/v2/tests/stop_instance_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_stop_instance_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -54,7 +54,7 @@ def test_stop_instance_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_stop_instance_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/update_instance_machine_type_test.py b/notebooks/v2/tests/update_instance_machine_type_test.py index c282c449198..a2bdb98b1cb 100644 --- a/notebooks/v2/tests/update_instance_machine_type_test.py +++ b/notebooks/v2/tests/update_instance_machine_type_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_update_instance_machine_type_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -59,7 +59,7 @@ def test_update_instance_machine_type_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_update_instance_machine_type_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/update_instance_metadata_test.py b/notebooks/v2/tests/update_instance_metadata_test.py index 9a20decb705..5ca1d578ead 100644 --- a/notebooks/v2/tests/update_instance_metadata_test.py +++ b/notebooks/v2/tests/update_instance_metadata_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_update_instance_metadata_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -67,7 +67,7 @@ def test_update_instance_metadata_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_update_instance_metadata_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/update_instance_shielded_config_test.py b/notebooks/v2/tests/update_instance_shielded_config_test.py index 413a2a75801..5818955aba3 100644 --- a/notebooks/v2/tests/update_instance_shielded_config_test.py +++ b/notebooks/v2/tests/update_instance_shielded_config_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_update_instance_shielded_config_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -74,7 +74,7 @@ def test_update_instance_shielded_config_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_update_instance_shielded_config_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], diff --git a/notebooks/v2/tests/upgrade_instance_test.py b/notebooks/v2/tests/upgrade_instance_test.py index 130d8c08fea..58aae0f3748 100644 --- a/notebooks/v2/tests/upgrade_instance_test.py +++ b/notebooks/v2/tests/upgrade_instance_test.py @@ -26,7 +26,7 @@ INSTANCE_NAME = f"projects/{PROJECT_ID}/locations/{LOCATION_ID}/instances/{INSTANCE_ID}" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_upgrade_instance_success( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], @@ -53,7 +53,7 @@ def test_upgrade_instance_success( assert captured.err == "" -@mock.patch("google.cloud.notebooks_v2.NotebookServiceClient") +@mock.patch.object(notebooks_v2, "NotebookServiceClient", autospec=True) def test_upgrade_instance_not_found( mock_client_cls: mock.MagicMock, capsys: pytest.CaptureFixture[str], From e2dcfea649ecda68a980fcd3f72f146c5b16f569 Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Thu, 5 Nov 2026 22:57:50 +0000 Subject: [PATCH 8/8] feat(notebooks): add explicit LRO timeouts and use dict assertions in complex v2 tests - Replace multi-line sequential attribute assertions in create_instance_hardened_test.py, create_instance_cmek_test.py, and create_instance_private_vpc_test.py with single-pass dictionary assertions via CreateInstanceRequest.to_dict(..., always_print_fields_with_no_presence=False, use_integers_for_enums=False). - Add explicit polling timeouts to operation.result(timeout=...) across all 18 LRO-based samples in notebooks/v2/ (timeout=300 for create/update/lifecycle operations and timeout=600 for upgrade_instance, rollback_instance, and diagnose_instance). --- notebooks/v2/create_instance.py | 2 +- notebooks/v2/create_instance_cmek.py | 2 +- notebooks/v2/create_instance_confidential.py | 2 +- .../v2/create_instance_custom_container.py | 2 +- notebooks/v2/create_instance_gpu.py | 2 +- notebooks/v2/create_instance_hardened.py | 2 +- notebooks/v2/create_instance_private_vpc.py | 2 +- notebooks/v2/create_instance_single_user.py | 2 +- notebooks/v2/delete_instance.py | 2 +- notebooks/v2/diagnose_instance.py | 2 +- notebooks/v2/reset_instance.py | 2 +- notebooks/v2/rollback_instance.py | 2 +- notebooks/v2/start_instance.py | 2 +- notebooks/v2/stop_instance.py | 2 +- .../v2/tests/create_instance_cmek_test.py | 36 +++++++-- .../v2/tests/create_instance_hardened_test.py | 81 ++++++++++++------- .../tests/create_instance_private_vpc_test.py | 36 ++++++--- notebooks/v2/update_instance_machine_type.py | 2 +- notebooks/v2/update_instance_metadata.py | 2 +- .../v2/update_instance_shielded_config.py | 2 +- notebooks/v2/upgrade_instance.py | 2 +- 21 files changed, 126 insertions(+), 63 deletions(-) diff --git a/notebooks/v2/create_instance.py b/notebooks/v2/create_instance.py index b3e3e65a3f7..96c7dd77ddb 100644 --- a/notebooks/v2/create_instance.py +++ b/notebooks/v2/create_instance.py @@ -46,7 +46,7 @@ def create_instance( print(f"Waiting for operation to complete: {operation.operation.name}") # The operation completes once the instance is accepted, before the # VM finishes provisioning, so the state is usually PROVISIONING here. - response = operation.result() + response = operation.result(timeout=300) print(f"Resource ID: {response.name}") print(f"Creation Date: {response.create_time}") print(f"Machine Type: {response.gce_setup.machine_type}") diff --git a/notebooks/v2/create_instance_cmek.py b/notebooks/v2/create_instance_cmek.py index 0203ae7c2f3..ae34ea45b64 100644 --- a/notebooks/v2/create_instance_cmek.py +++ b/notebooks/v2/create_instance_cmek.py @@ -72,7 +72,7 @@ def create_instance_cmek( print(f"Waiting for operation to complete: {operation.operation.name}") # The operation completes once the instance is accepted, before the # VM finishes provisioning, so the state is usually PROVISIONING here. - response = operation.result() + response = operation.result(timeout=300) boot_encryption = response.gce_setup.boot_disk.disk_encryption.name print(f"Resource ID: {response.name}") print(f"Creation Date: {response.create_time}") diff --git a/notebooks/v2/create_instance_confidential.py b/notebooks/v2/create_instance_confidential.py index ca35fb819a0..5a4ba70eac5 100644 --- a/notebooks/v2/create_instance_confidential.py +++ b/notebooks/v2/create_instance_confidential.py @@ -54,7 +54,7 @@ def create_instance_confidential( print(f"Waiting for operation to complete: {operation.operation.name}") # The operation completes once the instance is accepted, before the # VM finishes provisioning, so the state is usually PROVISIONING here. - response = operation.result() + response = operation.result(timeout=300) secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot print(f"Resource ID: {response.name}") print(f"Machine Type: {response.gce_setup.machine_type}") diff --git a/notebooks/v2/create_instance_custom_container.py b/notebooks/v2/create_instance_custom_container.py index 038d9e33e0c..0c6fcf738de 100644 --- a/notebooks/v2/create_instance_custom_container.py +++ b/notebooks/v2/create_instance_custom_container.py @@ -55,7 +55,7 @@ def create_instance_custom_container( print(f"Waiting for operation to complete: {operation.operation.name}") # The operation completes once the instance is accepted, before the # VM finishes provisioning, so the state is usually PROVISIONING here. - response = operation.result() + response = operation.result(timeout=300) repo = response.gce_setup.container_image.repository print(f"Resource ID: {response.name}") print(f"Container Repository: {repo}") diff --git a/notebooks/v2/create_instance_gpu.py b/notebooks/v2/create_instance_gpu.py index 6925bd75dc5..3b3ebd07950 100644 --- a/notebooks/v2/create_instance_gpu.py +++ b/notebooks/v2/create_instance_gpu.py @@ -57,7 +57,7 @@ def create_instance_gpu( print(f"Waiting for operation to complete: {operation.operation.name}") # The operation completes once the instance is accepted, before the # VM finishes provisioning, so the state is usually PROVISIONING here. - response = operation.result() + response = operation.result(timeout=300) print(f"Resource ID: {response.name}") print(f"Creation Date: {response.create_time}") print(f"Machine Type: {response.gce_setup.machine_type}") diff --git a/notebooks/v2/create_instance_hardened.py b/notebooks/v2/create_instance_hardened.py index 0cd6774694d..10e5b4a2c04 100644 --- a/notebooks/v2/create_instance_hardened.py +++ b/notebooks/v2/create_instance_hardened.py @@ -117,7 +117,7 @@ def create_instance_hardened( print(f"Waiting for operation to complete: {operation.operation.name}") # The operation completes once the instance is accepted, before the # VM finishes provisioning, so the state is usually PROVISIONING here. - response = operation.result() + response = operation.result(timeout=300) secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot print(f"Resource ID: {response.name}") print(f"Creation Date: {response.create_time}") diff --git a/notebooks/v2/create_instance_private_vpc.py b/notebooks/v2/create_instance_private_vpc.py index d077aff22fb..56233c60bf0 100644 --- a/notebooks/v2/create_instance_private_vpc.py +++ b/notebooks/v2/create_instance_private_vpc.py @@ -67,7 +67,7 @@ def create_instance_private_vpc( print(f"Waiting for operation to complete: {operation.operation.name}") # The operation completes once the instance is accepted, before the # VM finishes provisioning, so the state is usually PROVISIONING here. - response = operation.result() + response = operation.result(timeout=300) print(f"Resource ID: {response.name}") print(f"Creation Date: {response.create_time}") print(f"Disable Public IP: {response.gce_setup.disable_public_ip}") diff --git a/notebooks/v2/create_instance_single_user.py b/notebooks/v2/create_instance_single_user.py index 73b30402a9a..3b949dc5952 100644 --- a/notebooks/v2/create_instance_single_user.py +++ b/notebooks/v2/create_instance_single_user.py @@ -58,7 +58,7 @@ def create_instance_single_user( print(f"Waiting for operation to complete: {operation.operation.name}") # The operation completes once the instance is accepted, before the # VM finishes provisioning, so the state is usually PROVISIONING here. - response = operation.result() + response = operation.result(timeout=300) print(f"Resource ID: {response.name}") print(f"Creator: {response.creator}") print(f"State: {response.state.name}") diff --git a/notebooks/v2/delete_instance.py b/notebooks/v2/delete_instance.py index 5f5a7724561..82988dab1a7 100644 --- a/notebooks/v2/delete_instance.py +++ b/notebooks/v2/delete_instance.py @@ -38,7 +38,7 @@ def delete_instance( try: operation = client.delete_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") - operation.result() + operation.result(timeout=300) print(f"Resource ID: {name}") print(f"Location: {location_id}") print(f"Operation Done: {operation.done()}") diff --git a/notebooks/v2/diagnose_instance.py b/notebooks/v2/diagnose_instance.py index 2a6b33cdb71..e1720305d44 100644 --- a/notebooks/v2/diagnose_instance.py +++ b/notebooks/v2/diagnose_instance.py @@ -48,7 +48,7 @@ def diagnose_instance( try: operation = client.diagnose_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") - response = operation.result() + response = operation.result(timeout=600) print(f"Resource ID: {response.name}") print(f"Health State: {response.health_state.name}") print(f"State: {response.state.name}") diff --git a/notebooks/v2/reset_instance.py b/notebooks/v2/reset_instance.py index b24c88f2f37..056c54827c2 100644 --- a/notebooks/v2/reset_instance.py +++ b/notebooks/v2/reset_instance.py @@ -38,7 +38,7 @@ def reset_instance( try: operation = client.reset_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") - response = operation.result() + response = operation.result(timeout=300) print(f"Resource ID: {response.name}") print(f"State: {response.state.name}") except exceptions.NotFound as err: diff --git a/notebooks/v2/rollback_instance.py b/notebooks/v2/rollback_instance.py index eed58adde3e..c73a61d1ffa 100644 --- a/notebooks/v2/rollback_instance.py +++ b/notebooks/v2/rollback_instance.py @@ -47,7 +47,7 @@ def rollback_instance( try: operation = client.rollback_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") - response = operation.result() + response = operation.result(timeout=600) print(f"Resource ID: {response.name}") print(f"State: {response.state.name}") except exceptions.NotFound as err: diff --git a/notebooks/v2/start_instance.py b/notebooks/v2/start_instance.py index 880d5813a5b..177bb545ded 100644 --- a/notebooks/v2/start_instance.py +++ b/notebooks/v2/start_instance.py @@ -40,7 +40,7 @@ def start_instance( print(f"Waiting for operation to complete: {operation.operation.name}") # The operation completes before the VM finishes booting, so the # state is usually PROVISIONING here. - response = operation.result() + response = operation.result(timeout=300) print(f"Resource ID: {response.name}") print(f"Proxy URI: {response.proxy_uri}") print(f"State: {response.state.name}") diff --git a/notebooks/v2/stop_instance.py b/notebooks/v2/stop_instance.py index 1bac8a3bffd..e05415e3421 100644 --- a/notebooks/v2/stop_instance.py +++ b/notebooks/v2/stop_instance.py @@ -38,7 +38,7 @@ def stop_instance( try: operation = client.stop_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") - response = operation.result() + response = operation.result(timeout=300) print(f"Resource ID: {response.name}") print(f"State: {response.state.name}") except exceptions.NotFound as err: diff --git a/notebooks/v2/tests/create_instance_cmek_test.py b/notebooks/v2/tests/create_instance_cmek_test.py index c1188c5acec..ce28710498d 100644 --- a/notebooks/v2/tests/create_instance_cmek_test.py +++ b/notebooks/v2/tests/create_instance_cmek_test.py @@ -54,19 +54,39 @@ def test_create_instance_cmek_success( create_instance_cmek.create_instance_cmek(PROJECT_ID, LOCATION_ID, INSTANCE_ID) request = mock_client.create_instance.call_args.kwargs["request"] - assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" - assert request.instance_id == INSTANCE_ID - gce_setup = request.instance.gce_setup region_id = "-".join(LOCATION_ID.split("-")[:-1]) expected_kms_key = ( f"projects/{PROJECT_ID}/locations/{region_id}/" "keyRings/example-key-ring/cryptoKeys/example-key" ) - cmek = notebooks_v2.DiskEncryption.CMEK - assert gce_setup.boot_disk.disk_encryption == cmek - assert gce_setup.boot_disk.kms_key == expected_kms_key - assert gce_setup.data_disks[0].disk_encryption == cmek - assert gce_setup.data_disks[0].kms_key == expected_kms_key + actual_request = notebooks_v2.CreateInstanceRequest.to_dict( + request, + always_print_fields_with_no_presence=False, + use_integers_for_enums=False, + ) + assert actual_request == { + "parent": f"projects/{PROJECT_ID}/locations/{LOCATION_ID}", + "instance_id": INSTANCE_ID, + "instance": { + "gce_setup": { + "machine_type": "e2-standard-4", + "boot_disk": { + "disk_size_gb": "150", + "disk_type": "PD_SSD", + "disk_encryption": "CMEK", + "kms_key": expected_kms_key, + }, + "data_disks": [ + { + "disk_size_gb": "100", + "disk_type": "PD_SSD", + "disk_encryption": "CMEK", + "kms_key": expected_kms_key, + } + ], + }, + }, + } captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out diff --git a/notebooks/v2/tests/create_instance_hardened_test.py b/notebooks/v2/tests/create_instance_hardened_test.py index de847e7fe49..61bcef0ae4d 100644 --- a/notebooks/v2/tests/create_instance_hardened_test.py +++ b/notebooks/v2/tests/create_instance_hardened_test.py @@ -62,39 +62,66 @@ def test_create_instance_hardened_success( ) request = mock_client.create_instance.call_args.kwargs["request"] - assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" - assert request.instance_id == INSTANCE_ID - gce_setup = request.instance.gce_setup region_id = "-".join(LOCATION_ID.split("-")[:-1]) - network_interface = gce_setup.network_interfaces[0] - shielded = gce_setup.shielded_instance_config - metadata = dict(gce_setup.metadata) expected_kms_key = ( f"projects/{PROJECT_ID}/locations/{region_id}/" "keyRings/example-key-ring/cryptoKeys/example-key" ) - assert gce_setup.disable_public_ip - assert network_interface.network == ( - f"projects/{PROJECT_ID}/global/networks/example-network" + actual_request = notebooks_v2.CreateInstanceRequest.to_dict( + request, + always_print_fields_with_no_presence=False, + use_integers_for_enums=False, ) - assert network_interface.subnet == ( - f"projects/{PROJECT_ID}/regions/{region_id}/subnetworks/example-subnet" - ) - assert gce_setup.service_accounts[0].email == custom_sa_email - assert list(request.instance.instance_owners) == ["example-user@example.com"] - assert shielded.enable_secure_boot - assert shielded.enable_vtpm - assert shielded.enable_integrity_monitoring - cmek = notebooks_v2.DiskEncryption.CMEK - assert gce_setup.boot_disk.disk_encryption == cmek - assert gce_setup.boot_disk.kms_key == expected_kms_key - assert gce_setup.data_disks[0].disk_encryption == cmek - assert gce_setup.data_disks[0].kms_key == expected_kms_key - assert metadata["notebook-disable-root"] == "true" - assert metadata["notebook-disable-downloads"] == "true" - assert metadata["notebook-disable-terminal"] == "true" - assert metadata["serial-port-enable"] == "false" - assert metadata["block-project-ssh-keys"] == "true" + assert actual_request == { + "parent": f"projects/{PROJECT_ID}/locations/{LOCATION_ID}", + "instance_id": INSTANCE_ID, + "instance": { + "instance_owners": ["example-user@example.com"], + "gce_setup": { + "machine_type": "e2-standard-4", + "disable_public_ip": True, + "service_accounts": [{"email": custom_sa_email}], + "boot_disk": { + "disk_size_gb": "150", + "disk_type": "PD_SSD", + "disk_encryption": "CMEK", + "kms_key": expected_kms_key, + }, + "data_disks": [ + { + "disk_size_gb": "100", + "disk_type": "PD_SSD", + "disk_encryption": "CMEK", + "kms_key": expected_kms_key, + } + ], + "network_interfaces": [ + { + "network": ( + f"projects/{PROJECT_ID}/global/networks/example-network" + ), + "subnet": ( + f"projects/{PROJECT_ID}/regions/{region_id}/" + "subnetworks/example-subnet" + ), + "nic_type": "GVNIC", + } + ], + "shielded_instance_config": { + "enable_secure_boot": True, + "enable_vtpm": True, + "enable_integrity_monitoring": True, + }, + "metadata": { + "notebook-disable-root": "true", + "notebook-disable-downloads": "true", + "notebook-disable-terminal": "true", + "serial-port-enable": "false", + "block-project-ssh-keys": "true", + }, + }, + }, + } captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out diff --git a/notebooks/v2/tests/create_instance_private_vpc_test.py b/notebooks/v2/tests/create_instance_private_vpc_test.py index 325275a7867..5de23bd2277 100644 --- a/notebooks/v2/tests/create_instance_private_vpc_test.py +++ b/notebooks/v2/tests/create_instance_private_vpc_test.py @@ -51,18 +51,34 @@ def test_create_instance_private_vpc_success( ) request = mock_client.create_instance.call_args.kwargs["request"] - assert request.parent == f"projects/{PROJECT_ID}/locations/{LOCATION_ID}" - assert request.instance_id == INSTANCE_ID - gce_setup = request.instance.gce_setup region_id = "-".join(LOCATION_ID.split("-")[:-1]) - network_interface = gce_setup.network_interfaces[0] - assert gce_setup.disable_public_ip - assert network_interface.network == ( - f"projects/{PROJECT_ID}/global/networks/example-network" - ) - assert network_interface.subnet == ( - f"projects/{PROJECT_ID}/regions/{region_id}/subnetworks/example-subnet" + actual_request = notebooks_v2.CreateInstanceRequest.to_dict( + request, + always_print_fields_with_no_presence=False, + use_integers_for_enums=False, ) + assert actual_request == { + "parent": f"projects/{PROJECT_ID}/locations/{LOCATION_ID}", + "instance_id": INSTANCE_ID, + "instance": { + "gce_setup": { + "machine_type": "e2-standard-4", + "disable_public_ip": True, + "network_interfaces": [ + { + "network": ( + f"projects/{PROJECT_ID}/global/networks/example-network" + ), + "subnet": ( + f"projects/{PROJECT_ID}/regions/{region_id}/" + "subnetworks/example-subnet" + ), + "nic_type": "GVNIC", + } + ], + }, + }, + } captured = capsys.readouterr() assert f"Resource ID: {INSTANCE_NAME}" in captured.out diff --git a/notebooks/v2/update_instance_machine_type.py b/notebooks/v2/update_instance_machine_type.py index b133cde593a..b481ae616ff 100644 --- a/notebooks/v2/update_instance_machine_type.py +++ b/notebooks/v2/update_instance_machine_type.py @@ -47,7 +47,7 @@ def update_instance_machine_type( try: operation = client.update_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") - response = operation.result() + response = operation.result(timeout=300) print(f"Resource ID: {response.name}") print(f"Machine Type: {response.gce_setup.machine_type}") print(f"State: {response.state.name}") diff --git a/notebooks/v2/update_instance_metadata.py b/notebooks/v2/update_instance_metadata.py index e6e1eac804b..ae2f01e9c08 100644 --- a/notebooks/v2/update_instance_metadata.py +++ b/notebooks/v2/update_instance_metadata.py @@ -47,7 +47,7 @@ def update_instance_metadata( try: operation = client.update_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") - response = operation.result() + response = operation.result(timeout=300) # The service merges the update into metadata it manages itself, such # as the image version, so only the updated key changes. metadata = response.gce_setup.metadata diff --git a/notebooks/v2/update_instance_shielded_config.py b/notebooks/v2/update_instance_shielded_config.py index a22f7d78a7f..a8d0cbc349c 100644 --- a/notebooks/v2/update_instance_shielded_config.py +++ b/notebooks/v2/update_instance_shielded_config.py @@ -56,7 +56,7 @@ def update_instance_shielded_config( try: operation = client.update_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") - response = operation.result() + response = operation.result(timeout=300) secure_boot = response.gce_setup.shielded_instance_config.enable_secure_boot print(f"Resource ID: {response.name}") print(f"Secure Boot Enabled: {secure_boot}") diff --git a/notebooks/v2/upgrade_instance.py b/notebooks/v2/upgrade_instance.py index 7d4edd93daa..654d03cf753 100644 --- a/notebooks/v2/upgrade_instance.py +++ b/notebooks/v2/upgrade_instance.py @@ -38,7 +38,7 @@ def upgrade_instance( try: operation = client.upgrade_instance(request=request) print(f"Waiting for operation to complete: {operation.operation.name}") - response = operation.result() + response = operation.result(timeout=600) print(f"Resource ID: {response.name}") print(f"State: {response.state.name}") except exceptions.NotFound as err: