tests(containers): add e2e test for MySQL with user auth (#1405)

Author: jiridanek

tests/containers/conftest.py

@@ -15,9 +15,10 @@
 import testcontainers.core.docker_client
 
 from tests.containers import docker_utils, skopeo_utils, utils
+from tests.containers.kubernetes_utils import TestFrame
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Callable, Generator
 
     from pytest import ExitCode, Metafunc, Parser, Session
 
@@ -108,6 +109,12 @@ def skip_if_not_rocm_image(image: str) -> Image:
     return image_metadata
 
 
+@pytest.fixture(scope="function")
+def tf() -> Generator[TestFrame[Any], None, None]:
+    with TestFrame() as tf:
+        yield tf
+
+
 # https://docs.pytest.org/en/stable/how-to/fixtures.html#parametrizing-fixtures
 # indirect parametrization https://stackoverflow.com/questions/18011902/how-to-pass-a-parameter-to-a-fixture-function-in-pytest
 @pytest.fixture(scope="session")
@@ -172,6 +179,17 @@ def jupyterlab_trustyai_image(jupyterlab_image: Image) -> Image:
     return jupyterlab_image
 
 
+@pytest.fixture(scope="function")
+def datascience_image(image: str) -> Image:
+    image_metadata = get_image_metadata(image)
+    if "-minimal-" in image_metadata.labels["name"]:
+        pytest.skip(
+            f"Image {image_metadata.name} is not datascience image because it has '-minimal-' in {image_metadata.labels['name']=}'"
+        )
+
+    return image_metadata
+
+
 @pytest.fixture(scope="function")
 def rstudio_image(image: str) -> Image:
     image_metadata = skip_if_not_workbench_image(image)

tests/containers/docker_utils.py

@@ -3,6 +3,7 @@
 import io
 import logging
 import os.path
+import socket as pysocket
 import sys
 import tarfile
 import time
@@ -163,6 +164,105 @@ def communicate(self, line_prefix=b"") -> int:
         return self.poll()
 
 
+def container_exec_with_stdin(
+    container: Container,
+    cmd: str | list[str],
+    stdin_data: str | bytes,
+    read_timeout: float = 1.0,
+) -> tuple[int, bytes]:
+    """
+    Executes a command in a container, writing stdin_data to its stdin.
+
+    :param container: The container to execute the command in.
+    :param cmd: The command to execute.
+    :param stdin_data: The string or bytes to send to the command's stdin.
+    :return: A tuple of (exit_code, output_bytes).
+    """
+    if isinstance(stdin_data, str):
+        stdin_data = stdin_data.encode("utf-8")
+
+    # Using the low-level API for precise control over the socket.
+    exec_id = container.client.api.exec_create(
+        container=container.id,
+        cmd=cmd,
+        stdin=True,
+        stdout=True,
+        stderr=True,
+        tty=True,
+    )
+
+    # When using a podman client, exec_start(socket=True) returns a file-like
+    # object (a wrapper around SocketIO), not a raw socket. We must use
+    # file-like methods (write, read) instead of raw socket methods.
+    stream = container.client.api.exec_start(exec_id, socket=True, tty=True)
+
+    # The stream object can be a raw socket or a file-like wrapper which might
+    # be incorrectly marked as read-only. We need to find the underlying raw
+    # socket to reliably write to stdin.
+    raw_sock = None
+    if isinstance(stream, pysocket.socket):
+        raw_sock = stream
+    else:
+        # Try to unwrap a file-like object (e.g., BufferedReader -> SocketIO -> socket)
+        raw_io = getattr(stream, "raw", stream)
+        if hasattr(raw_io, "_sock"):
+            raw_sock = raw_io._sock
+
+    if raw_sock:
+        raw_sock.sendall(stdin_data)
+    else:
+        # Fallback to stream.write() if no raw socket found. This may fail.
+        try:
+            stream.write(stdin_data)
+            stream.flush()
+        except (OSError, io.UnsupportedOperation) as e:
+            raise OSError(f"Could not write to container exec stdin using stream of type {type(stream)}") from e
+
+    # Shut down the write-half of the connection to signal EOF to the process.
+    try:
+        if raw_sock:
+            raw_sock.shutdown(pysocket.SHUT_WR)
+        else:
+            # Fallback for stream objects that have a shutdown method.
+            raw_io = getattr(stream, "raw", stream)
+            if hasattr(raw_io, "_sock"):
+                raw_io._sock.shutdown(pysocket.SHUT_WR)
+            else:
+                stream.shutdown(pysocket.SHUT_WR)
+    except (OSError, AttributeError):
+        # This is expected if the remote process closes the connection first.
+        pass
+
+    if raw_sock:
+        # If we unwrapped and used the raw socket, we must continue using it
+        # for reading to avoid state inconsistencies with the wrapper object.
+        output_chunks = []
+        while True:
+            # we set the timeout in order not to be blocked afterwards with blocking read
+            raw_sock.settimeout(read_timeout)
+            # Reading in a loop is the standard way to consume a socket's content.
+            try:
+                chunk = raw_sock.recv(4096)
+            except TimeoutError:
+                break
+            if not chunk:
+                # An empty chunk signifies that the remote end has closed the connection.
+                break
+            output_chunks.append(chunk)
+        output = b"".join(output_chunks)
+        raw_sock.close()
+    else:
+        # Fallback to stream.read() if we couldn't get a raw socket.
+        # This may hang if the shutdown logic above also failed.
+        output = stream.read()
+        stream.close()
+
+    # Get the exit code of the process.
+    exit_code = container.client.api.exec_inspect(exec_id)["ExitCode"]
+
+    return exit_code, output
+
+
 def get_socket_path(client: docker.client.DockerClient) -> str:
     """Determine the local socket path.
     This works even when `podman machine` with its own host-mounts is involved

tests/containers/kubernetes_utils.py

@@ -90,6 +90,8 @@ def test_get_username(self):
 
 
 class TestFrame[S]:
+    __test__ = False
+
     def __init__(self):
         self.stack: list[tuple[S, Callable[[S], None] | None]] = []
 
@@ -102,6 +104,7 @@ def defer_resource[T: ocp_resources.resource.Resource](
 
     def defer[T](self, obj: T, destructor: Callable[[T], None] | None = None) -> T:
         self.stack.append((obj, destructor))
+        return obj
 
     def destroy(self, wait=False):
         while self.stack:

tests/containers/workbenches/jupyterlab/jupyterlab_datascience_test.py

@@ -2,12 +2,21 @@
 
 import pathlib
 import tempfile
+import typing
 
 import allure
+import pytest
+import testcontainers.core.network
+from testcontainers.core.waiting_utils import wait_for_logs
+from testcontainers.mysql import MySqlContainer
 
 from tests.containers import conftest, docker_utils
 from tests.containers.workbenches.workbench_image_test import WorkbenchContainer
 
+if typing.TYPE_CHECKING:
+    from tests.containers.conftest import Image
+    from tests.containers.kubernetes_utils import TestFrame
+
 
 class TestJupyterLabDatascienceImage:
     """Tests for JupyterLab Workbench images in this repository that are not -minimal-."""
@@ -70,3 +79,112 @@ def test_sklearn_smoke(self, jupyterlab_datascience_image: conftest.Image) -> No
 
         finally:
             docker_utils.NotebookContainer(container).stop(timeout=0)
+
+    @allure.description("Check that mysql client functionality is working with SASL plain auth.")
+    def test_mysql_connection(self, tf: TestFrame, datascience_image: Image, subtests):
+        network = testcontainers.core.network.Network()
+        tf.defer(network.create())
+
+        mysql_container = (
+            MySqlContainer("docker.io/library/mysql:9.3.0").with_network(network).with_network_aliases("mysql")
+        )
+        tf.defer(mysql_container.start())
+
+        try:
+            wait_for_logs(mysql_container, r"mysqld: ready for connections.", timeout=30)
+        except TimeoutError:
+            print("Container is not ready.")
+            print(mysql_container.get_wrapped_container().logs(stdout=True, stderr=True))
+            raise
+        print("Container is ready. Setting up test user...")
+
+        host = "mysql"
+        port = 3306
+
+        # language=Python
+        setup_mysql_user = f"""
+import mysql.connector
+
+conn = mysql.connector.connect(
+    user='root',
+    password='{mysql_container.root_password}',
+    host = "{host}",
+    port = {port},
+)
+cursor = conn.cursor()
+print("Creating test users...")
+
+cursor.execute(
+# language=mysql
+'''
+CREATE USER 'clearpassuser'@'%' IDENTIFIED WITH caching_sha2_password BY 'clearpassword';
+GRANT ALL PRIVILEGES ON *.* TO 'clearpassuser'@'%';
+
+FLUSH PRIVILEGES;
+''')
+cursor.close()
+conn.close()
+print("Test users created successfully.")
+"""
+
+        # language=Python
+        clearpassuser = f"""
+import mysql.connector
+
+try:
+    cnx = mysql.connector.connect(
+        user='clearpassuser',
+        password='clearpassword',
+        host='{host}',
+        port={port},
+        auth_plugin='mysql_clear_password',
+    )
+    cursor = cnx.cursor()
+    cursor.execute("SELECT 1")
+    result = cursor.fetchone()
+    if result == (1,):
+        print("MySQL connection successful!")
+    else:
+        print("MySQL connection failed!")
+    cnx.close()
+except Exception as e:
+    print(f"An error occurred: {{e}}")
+    raise
+"""
+
+        container = WorkbenchContainer(image=datascience_image.name, user=4321, group_add=[0])
+        (container.with_network(network).with_command("/bin/sh -c 'sleep infinity'"))
+        try:
+            container.start(wait_for_readiness=False)
+
+            # RHOAIENG-140: code-server image users are expected to install their own db clients
+            if "-code-server-" in datascience_image.labels["name"]:
+                exit_code, output = container.exec(["python", "-m", "pip", "install", "mysql-connector-python==9.3.0"])
+                output_str = output.decode()
+                print(output_str)
+
+                assert exit_code == 0, f"Failed to install mysql-connector-python: {output_str}"
+            elif "-rstudio-" in datascience_image.labels["name"]:
+                pytest.skip(
+                    f"Image {datascience_image.name} does have -rstudio- in {datascience_image.labels['name']=}'"
+                )
+
+            with subtests.test("Setting the user..."):
+                exit_code, output = container.exec(["python", "-c", setup_mysql_user])
+                output_str = output.decode()
+
+                print(output_str)
+
+                assert "Test users created successfully." in output_str
+                assert exit_code == 0
+
+            with subtests.test("Checking the output of the clearpassuser script..."):
+                exit_code, output = container.exec(["python", "-c", clearpassuser])
+                output_str = output.decode()
+
+                print(output_str)
+
+                assert "MySQL connection successful!" in output_str
+                assert exit_code == 0
+        finally:
+            docker_utils.NotebookContainer(container).stop(timeout=0)