Add presubmit integration testing

README.md

@@ -24,7 +24,7 @@ If you are running the client outside of Google Cloud, you must set following en
 * GOOGLE_CLOUD_PROJECT - The Google Cloud project you use to run Spark workloads
 * GOOGLE_CLOUD_REGION - The Compute Engine [region](https://cloud.google.com/compute/docs/regions-zones#available) where you run the Spark workload.
 * GOOGLE_APPLICATION_CREDENTIALS - Your [Application Credentials](https://cloud.google.com/docs/authentication/provide-credentials-adc)
-* DATAPROC_SPARK_CONNECT_SESSION_DEFAULT_CONFIG (Optional) - The config location, such as `tests/integration/resources/session.textproto`
+* DATAPROC_SPARK_CONNECT_SESSION_DEFAULT_CONFIG (Optional) - The config location, such as `tests/integration/resources/session_user.textproto`
 
 ## Usage
 

cloudbuild/Dockerfile

@@ -1,5 +1,13 @@
 FROM python:3.10-bookworm
 
+RUN \
+  echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | \
+  tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
+  curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | \
+  apt-key --keyring /usr/share/keyrings/cloud.google.gpg  add - && \
+  apt-get update -y && \
+  apt-get install google-cloud-cli jq -y
+
 WORKDIR /opt/tests
 # NOTE: We copy the requirements files explicitly _before_ installing
 # dependencies to allow layer caching. Afterward, we copy the rest of the build
@@ -11,4 +19,7 @@ RUN python3 -m pip install -U pip
 RUN python3 -m pip install --no-cache-dir -r requirements-dev.txt -r requirements-test.txt
 COPY . .
 RUN python setup.py sdist bdist_wheel egg_info
+# TODO: Figure out how to get the .egg-info/ directory included in the output
+# image. Something is suppressing it at the moment, so we have to `pip install`
+# this package to allow tests to run.
 RUN pip install .

cloudbuild/cloudbuild.yaml

@@ -3,11 +3,33 @@ steps:
   # distribution artifacts.
   - name: 'gcr.io/cloud-builders/docker'
     id: 'build-container-image'
-    args: ['build', '--tag=gcr.io/${PROJECT_ID}/google-spark-connect/google-spark-connect-presubmit:${BUILD_ID}', -f, 'cloudbuild/Dockerfile', '.']
+    args: ['build', '--tag=gcr.io/${PROJECT_ID}/google-spark-connect/google-spark-connect-presubmit', -f, 'cloudbuild/Dockerfile', '.']
   # Run all unit tests
-  - name: 'gcr.io/${PROJECT_ID}/google-spark-connect/google-spark-connect-presubmit:${BUILD_ID}'
+  - name: 'gcr.io/${PROJECT_ID}/google-spark-connect/google-spark-connect-presubmit'
     id: 'run-unit-tests'
     waitFor: ['build-container-image']
-    entrypoint: 'pytest'
-    args: ['-n', 'auto', 'tests/unit']
-timeout: 600s
+    allowFailure: true
+    entrypoint: '/opt/tests/cloudbuild/run-unit-tests.sh'
+  # Run all integration tests. These exercise the actual API endpoints and
+  # require credentials.
+  - name: 'gcr.io/${PROJECT_ID}/google-spark-connect/google-spark-connect-presubmit'
+    id: 'run-integration-tests'
+    waitFor: ['build-container-image']
+    allowFailure: true
+    entrypoint: '/opt/tests/cloudbuild/run-integration-tests.sh'
+    env:
+      - 'GOOGLE_CLOUD_SUBNET=${_GOOGLE_CLOUD_SUBNET}'
+      - 'GOOGLE_CLOUD_PROJECT=${_GOOGLE_CLOUD_PROJECT}'
+      - 'GOOGLE_CLOUD_REGION=${_GOOGLE_CLOUD_REGION}'
+      - 'GOOGLE_CLOUD_SERVICE_ACCOUNT=${_GOOGLE_CLOUD_SERVICE_ACCOUNT}'
+      - 'TEST_SUPPRESS_END_USER_CREDENTIALS=1'
+  - name: 'gcr.io/${PROJECT_ID}/google-spark-connect/google-spark-connect-presubmit'
+    id: 'verify-status'
+    waitFor: ['run-unit-tests', 'run-integration-tests']
+    entrypoint: '/opt/tests/cloudbuild/verify-status.sh'
+    env:
+      - 'BUILD_ID=${BUILD_ID}'
+      - 'GOOGLE_CLOUD_PROJECT=${PROJECT_ID}'
+timeout: 1800s
+options:
+    defaultLogsBucketBehavior: REGIONAL_USER_OWNED_BUCKET

cloudbuild/run-integration-tests.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+function main() {
+    run_tests && touch /workspace/integration-tests.SUCCESS
+}
+
+function run_tests() {
+    pytest -n 10 --log-cli-level=DEBUG tests/integration
+}
+
+main "$@"

cloudbuild/run-unit-tests.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+function main() {
+    run_tests && touch /workspace/unit-tests.SUCCESS
+}
+
+function run_tests() {
+    pytest -n 10 tests/unit
+}
+
+main "$@"

cloudbuild/verify-status.sh

@@ -0,0 +1,59 @@
+#!/usr/bin/env bash
+
+# TODO: Consider rewriting all of these test drivers into a single python script
+# to share names and logic.
+
+set -euo pipefail
+
+readonly TESTS=(
+    'unit-tests'
+    'integration-tests'
+)
+
+function main() {
+    #echo "/workspace contents:"
+    #ls -R /workspace
+    local failed_tests=()
+    local sentinel=""
+    for test in "${TESTS[@]}" ; do
+        # sentinel="/tmp/${test}.SUCCESS"
+        sentinel="/workspace/${test}.SUCCESS"
+        if [ ! -f "$sentinel" ] ; then
+            failed_tests+=("$test")
+        fi
+    done
+    if [ "${#failed_tests[@]}" -gt 0 ] ; then
+        echo "failed tests: ${failed_tests[@]}"
+        return 1
+    fi
+    # echo ""
+    # echo "BUILD RESULTS:"
+    # describe_build
+    # describe_build | verify_statuses
+    # jq_program
+}
+
+function describe_build() {
+    gcloud builds describe \
+        --project "$GOOGLE_CLOUD_PROJECT" \
+        --format=json \
+        "$BUILD_ID"
+}
+
+function verify_statuses() {
+    jq -f <(jq_program)
+}
+
+function jq_program() {
+    cat <<'EOF'
+reduce
+    (.steps[] | select(.allowFailure)) as $step
+    (
+        [];
+        if $step.status != "SUCCESS" then . + [$step.id] else . end
+    ) |
+    if length > 0 then error("the following steps failed: " + (join(", "))) else empty end
+EOF
+}
+
+main "$@"

google/cloud/spark_connect/session.py

@@ -172,6 +172,7 @@ def __create_spark_connect_session_from_s8s(
             return session
 
         def __create(self) -> "SparkSession":
+            print("CREATING SPARK SESSION", flush=True)
             with self._lock:
 
                 if self._options.get("spark.remote", False):
@@ -208,6 +209,8 @@ def __create(self) -> "SparkSession":
                     f"projects/{self._project_id}/locations/{self._region}"
                 )
 
+                print("FULL SESSION REQUEST:", session_request, flush=True)
+
                 logger.debug("Creating serverless session")
                 GoogleSparkSession._active_s8s_session_id = session_id
                 s8s_creation_start_time = time.time()

tests/integration/resources/session_service_account.textproto

@@ -0,0 +1,12 @@
+environment_config: {
+  execution_config: {
+    subnetwork_uri: "subnet-placeholder"
+    service_account: "service-account-placeholder"
+    authentication_config: {
+        user_workload_authentication_type: SERVICE_ACCOUNT
+    }
+  }
+}
+runtime_config: {
+  version: "2.2"
+}

tests/integration/resources/session.textproto → tests/integration/resources/session_user.textproto

@@ -1,6 +1,9 @@
 environment_config: {
   execution_config: {
     subnetwork_uri: "subnet-placeholder"
+    authentication_config: {
+        user_workload_authentication_type: END_USER_CREDENTIALS
+    }
   }
 }
 runtime_config: {

tests/integration/test_session.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import datetime
+import enum
 import os
 import tempfile
 import uuid
@@ -40,29 +41,79 @@
 _SERVICE_ACCOUNT_KEY_FILE_ = "service_account_key.json"
 
 
+class AuthType(enum.Enum):
+    SERVICE_ACCOUNT = "SERVICE_ACCOUNT"
+    END_USER_CREDENTIALS = "END_USER_CREDENTIALS"
+
+
+def get_auth_types():
+    service_account = os.environ.get("GOOGLE_CLOUD_SERVICE_ACCOUNT")
+    suppress_end_user_creds = os.environ.get(
+        "TEST_SUPPRESS_END_USER_CREDENTIALS"
+    )
+    if service_account:
+        yield AuthType.SERVICE_ACCOUNT
+    if not suppress_end_user_creds or suppress_end_user_creds == "0":
+        yield AuthType.END_USER_CREDENTIALS
+
+
+def get_config_path(auth_type):
+    resources_dir = os.path.join(os.path.dirname(__file__), "resources")
+    match auth_type:
+        case AuthType.SERVICE_ACCOUNT:
+            return os.path.join(
+                resources_dir, "session_service_account.textproto"
+            )
+        case AuthType.END_USER_CREDENTIALS:
+            return os.path.join(resources_dir, "session_user.textproto")
+        case _:
+            raise Exception(f"unknown auth_type: {auth_type}")
+
+
 @pytest.fixture(params=["2.2", "3.0"])
 def image_version(request):
     return request.param
 
 
 @pytest.fixture
 def test_project():
-    return os.environ.get("GOOGLE_CLOUD_PROJECT")
+    project = os.environ.get("GOOGLE_CLOUD_PROJECT")
+    if project is None:
+        raise Exception(
+            "must set a project through GOOGLE_CLOUD_PROJECT environment variable"
+        )
+    return project
 
 
-@pytest.fixture
+@pytest.fixture(params=get_auth_types())
 def auth_type(request):
-    return getattr(request, "param", "SYSTEM_SERVICE_ACCOUNT")
+    return request.param
 
 
 @pytest.fixture
 def test_region():
-    return os.environ.get("GOOGLE_CLOUD_REGION")
+    region = os.environ.get("GOOGLE_CLOUD_REGION")
+    if region is None:
+        raise Exception(
+            "must set region through GOOGLE_CLOUD_REGION environment variable"
+        )
+    return region
 
 
 @pytest.fixture
 def test_subnet():
-    return os.environ.get("GOOGLE_CLOUD_SUBNET")
+    subnet = os.environ.get("GOOGLE_CLOUD_SUBNET")
+    if subnet is None:
+        raise Exception(
+            "must set subnet through GOOGLE_CLOUD_SUBNET environment variable"
+        )
+    return subnet
+
+
+@pytest.fixture
+def test_service_account():
+    # The service account may be empty, in which case we skip SA-based testing.
+    return os.environ.get("GOOGLE_CLOUD_SERVICE_ACCOUNT")
 
 
 @pytest.fixture
@@ -72,22 +123,29 @@ def test_subnetwork_uri(test_project, test_region, test_subnet):
 
 @pytest.fixture
 def default_config(
-    auth_type, image_version, test_project, test_region, test_subnetwork_uri
+    auth_type,
+    test_service_account,
+    image_version,
+    test_project,
+    test_region,
+    test_subnetwork_uri,
 ):
-    resources_dir = os.path.join(os.path.dirname(__file__), "resources")
-    template_file = os.path.join(resources_dir, "session.textproto")
+    template_file = get_config_path(auth_type)
     with open(template_file) as f:
         template = f.read()
-        contents = (
-            template.replace("2.2", image_version)
-            .replace("subnet-placeholder", test_subnetwork_uri)
-            .replace("SYSTEM_SERVICE_ACCOUNT", auth_type)
+        contents = template.replace("2.2", image_version).replace(
+            "subnet-placeholder", test_subnetwork_uri
         )
-        with tempfile.NamedTemporaryFile(delete=False) as t:
-            t.write(contents.encode("utf-8"))
-            t.close()
-            yield t.name
-            os.remove(t.name)
+        if auth_type == AuthType.SERVICE_ACCOUNT:
+            contents = contents.replace(
+                "service-account-placeholder", test_service_account
+            )
+    print("CONFIG CONTENTS:", contents)
+    with tempfile.NamedTemporaryFile(delete=False) as t:
+        t.write(contents.encode("utf-8"))
+        t.close()
+        yield t.name
+        os.remove(t.name)
 
 
 @pytest.fixture
@@ -128,6 +186,7 @@ def session_template_controller_client(test_client_options):
 
 @pytest.fixture
 def connect_session(test_project, test_region, os_environment):
+    print("CREATING SESSION (TEST)", flush=True)
     return GoogleSparkSession.builder.getOrCreate()
 
 
@@ -136,12 +195,16 @@ def session_name(test_project, test_region, connect_session):
     return f"projects/{test_project}/locations/{test_region}/sessions/{GoogleSparkSession._active_s8s_session_id}"
 
 
-@pytest.mark.parametrize("auth_type", ["END_USER_CREDENTIALS"], indirect=True)
+# @pytest.mark.parametrize("auth_type", ["END_USER_CREDENTIALS"], indirect=True)
 def test_create_spark_session_with_default_notebook_behavior(
-    auth_type, connect_session, session_name, session_controller_client
+    connect_session,
+    session_name,
+    session_controller_client,
 ):
+    # print("auth type parameter:", auth_type)
     get_session_request = GetSessionRequest()
     get_session_request.name = session_name
+    print("GET SESSION REQUEST (TEST):", get_session_request, flush=True)
     session = session_controller_client.get_session(get_session_request)
     assert session.state == Session.State.ACTIVE
 
@@ -307,13 +370,28 @@ def session_template_name(
 
 
 def test_create_spark_session_with_session_template_and_user_provided_dataproc_config(
+    auth_type,
+    test_service_account,
     image_version,
     test_project,
     test_region,
     session_template_name,
     session_controller_client,
 ):
     dataproc_config = Session()
+    match auth_type:
+        case AuthType.END_USER_CREDENTIALS:
+            # This is the default
+            pass
+        case AuthType.SERVICE_ACCOUNT:
+            dataproc_config.environment_config.execution_config.service_account = (
+                test_service_account
+            )
+            dataproc_config.environment_config.execution_config.authentication_config.user_workload_authentication_type = (
+                "SERVICE_ACCOUNT"
+            )
+        case _:
+            raise Exception(f"unknown auth_type: {auth_type}")
     dataproc_config.environment_config.execution_config.ttl = {"seconds": 64800}
     dataproc_config.session_template = session_template_name
     connect_session = (