fixit: Clean up Python sample at cloud-sql/postgres/client-side-encry… (GoogleCloudPlatform#10027)

m-strzelczyk · web-flow · commit 68dd06451242 · 2023-05-25T13:12:00.000-07:00
* fixit: Clean up Python sample at cloud-sql/postgres/client-side-encryption

* Removing capsys from tests
diff --git a/cloud-sql/postgres/client-side-encryption/snippets/cloud_kms_env_aead.py b/cloud-sql/postgres/client-side-encryption/snippets/cloud_kms_env_aead.py
@@ -22,9 +22,10 @@
 logger = logging.getLogger(__name__)
 
 
-def init_tink_env_aead(
-        key_uri: str,
-        credentials: str) -> tink.aead.KmsEnvelopeAead:
+def init_tink_env_aead(key_uri: str, credentials: str) -> tink.aead.KmsEnvelopeAead:
+    """
+    Initiates the Envelope AEAD object using the KMS credentials.
+    """
     aead.register()
 
     try:
diff --git a/cloud-sql/postgres/client-side-encryption/snippets/cloud_kms_env_aead_test.py b/cloud-sql/postgres/client-side-encryption/snippets/cloud_kms_env_aead_test.py
@@ -26,15 +26,14 @@ def setup() -> str:
     yield kms_uri
 
 
-def test_cloud_kms_env_aead(
-        capsys: pytest.CaptureFixture, kms_uri: str) -> None:
+def test_cloud_kms_env_aead(kms_uri: str) -> None:
     credentials = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", None)
     if credentials is None:
         raise Exception(
-            "Environment variable GOOGLE_APPLICATION_CREDENTIALS is not set")
+            "Environment variable GOOGLE_APPLICATION_CREDENTIALS is not set"
+        )
 
     # Create env_aead primitive
-    init_tink_env_aead(kms_uri, credentials)
+    envelope = init_tink_env_aead(kms_uri, credentials)
 
-    captured = capsys.readouterr().out
-    assert f"Created envelope AEAD Primitive using KMS URI: {kms_uri}" in captured
+    assert envelope.key_template == kms_uri
diff --git a/cloud-sql/postgres/client-side-encryption/snippets/cloud_sql_connection_pool.py b/cloud-sql/postgres/client-side-encryption/snippets/cloud_sql_connection_pool.py
@@ -19,6 +19,9 @@
 def init_tcp_connection_engine(
     db_user: str, db_pass: str, db_name: str, db_host: str
 ) -> sqlalchemy.engine.base.Engine:
+    """
+    Creates a connection to the database using tcp socket.
+    """
     # Remember - storing secrets in plaintext is potentially unsafe. Consider using
     # something like https://cloud.google.com/secret-manager/docs/overview to help keep
     # secrets secret.
@@ -50,6 +53,9 @@ def init_unix_connection_engine(
     instance_connection_name: str,
     db_socket_dir: str,
 ) -> sqlalchemy.engine.base.Engine:
+    """
+    Creates a connection to the database using unix socket.
+    """
     # Remember - storing secrets in plaintext is potentially unsafe. Consider using
     # something like https://cloud.google.com/secret-manager/docs/overview to help keep
     # secrets secret.
@@ -64,9 +70,9 @@ def init_unix_connection_engine(
             database=db_name,  # e.g. "my-database-name"
             query={
                 "unix_sock": "{}/{}/.s.PGSQL.5432".format(
-                    db_socket_dir,  # e.g. "/cloudsql"
-                    instance_connection_name)  # i.e "<PROJECT-NAME>:<INSTANCE-REGION>:<INSTANCE-NAME>"
-            }
+                    db_socket_dir, instance_connection_name  # e.g. "/cloudsql"
+                )  # i.e "<PROJECT-NAME>:<INSTANCE-REGION>:<INSTANCE-NAME>"
+            },
         ),
     )
     print("Created Unix socket connection pool")
@@ -82,6 +88,7 @@ def init_db(
     db_socket_dir: str = None,
     db_host: str = None,
 ) -> sqlalchemy.engine.base.Engine:
+    """Starts a connection to the database and creates voting table if it doesn't exist."""
 
     if db_host:
         db = init_tcp_connection_engine(db_user, db_pass, db_name, db_host)
diff --git a/cloud-sql/postgres/client-side-encryption/snippets/cloud_sql_connection_pool_test.py b/cloud-sql/postgres/client-side-encryption/snippets/cloud_sql_connection_pool_test.py
@@ -18,11 +18,12 @@
 import uuid
 
 import pytest
+import sqlalchemy
 
 from snippets.cloud_sql_connection_pool import (
     init_db,
     init_tcp_connection_engine,
-    init_unix_connection_engine
+    init_unix_connection_engine,
 )
 
 
@@ -46,50 +47,46 @@ def setup() -> dict[str, str]:
         yield conn_vars
 
 
-def test_init_tcp_connection_engine(
-        capsys: pytest.CaptureFixture,
-        conn_vars: dict[str, str]) -> None:
-
-    init_tcp_connection_engine(
+def test_init_tcp_connection_engine(conn_vars: dict[str, str]) -> None:
+    engine = init_tcp_connection_engine(
         db_user=conn_vars["db_user"],
         db_name=conn_vars["db_name"],
         db_pass=conn_vars["db_pass"],
         db_host=conn_vars["db_host"],
     )
 
-    captured = capsys.readouterr().out
-    assert "Created TCP connection pool" in captured
-
+    assert isinstance(engine, sqlalchemy.engine.base.Engine)
+    assert conn_vars["db_name"] in engine.url
 
-def test_init_unix_connection_engine(
-        capsys: pytest.CaptureFixture,
-        conn_vars: dict[str, str]) -> None:
 
-    init_unix_connection_engine(
+def test_init_unix_connection_engine(conn_vars: dict[str, str]) -> None:
+    engine = init_unix_connection_engine(
         db_user=conn_vars["db_user"],
         db_name=conn_vars["db_name"],
         db_pass=conn_vars["db_pass"],
         instance_connection_name=conn_vars["instance_conn_name"],
         db_socket_dir=conn_vars["db_socket_dir"],
     )
 
-    captured = capsys.readouterr().out
-    assert "Created Unix socket connection pool" in captured
-
+    assert isinstance(engine, sqlalchemy.engine.base.Engine)
+    assert conn_vars["db_name"] in engine.url
 
-def test_init_db(
-        capsys: pytest.CaptureFixture,
-        conn_vars: dict[str, str]) -> None:
 
+def test_init_db(conn_vars: dict[str, str]) -> None:
     table_name = f"votes_{uuid.uuid4().hex}"
 
-    init_db(
+    engine = init_db(
         db_user=conn_vars["db_user"],
         db_name=conn_vars["db_name"],
         db_pass=conn_vars["db_pass"],
         table_name=table_name,
         db_host=conn_vars["db_host"],
     )
 
-    captured = capsys.readouterr().out
-    assert f"Created table {table_name} in db {conn_vars['db_name']}" in captured
+    assert isinstance(engine, sqlalchemy.engine.base.Engine)
+
+    try:
+        with engine.connect() as conn:
+            conn.execute(f"SELECT count(*) FROM {table_name}").all()
+    except Exception as error:
+        pytest.fail(f"Database wasn't initialized properly: {error}")
diff --git a/cloud-sql/postgres/client-side-encryption/snippets/encrypt_and_insert_data.py b/cloud-sql/postgres/client-side-encryption/snippets/encrypt_and_insert_data.py
@@ -28,6 +28,9 @@
 
 
 def main() -> None:
+    """
+    Connects to the database, encrypts and inserts some data.
+    """
     db_user = os.environ["DB_USER"]  # e.g. "root", "postgres"
     db_pass = os.environ["DB_PASS"]  # e.g. "mysupersecretpassword"
     db_name = os.environ["DB_NAME"]  # e.g. "votes_db"
@@ -73,6 +76,10 @@ def encrypt_and_insert_data(
     team: str,
     email: str,
 ) -> None:
+    """
+    Inserts a vote into the database with email address previously encrypted using
+    a KmsEnvelopeAead object.
+    """
     time_cast = datetime.datetime.now(tz=datetime.timezone.utc)
     # Use the envelope AEAD primitive to encrypt the email, using the team name as
     # associated data. Encryption with associated data ensures authenticity
@@ -93,11 +100,7 @@ def encrypt_and_insert_data(
     # Using a with statement ensures that the connection is always released
     # back into the pool at the end of statement (even if an error occurs)
     with db.connect() as conn:
-        conn.execute(
-            stmt,
-            time_cast=time_cast,
-            team=team,
-            voter_email=encrypted_email)
+        conn.execute(stmt, time_cast=time_cast, team=team, voter_email=encrypted_email)
     print(f"Vote successfully cast for '{team}' at time {time_cast}!")
 
 
diff --git a/cloud-sql/postgres/client-side-encryption/snippets/encrypt_and_insert_data_test.py b/cloud-sql/postgres/client-side-encryption/snippets/encrypt_and_insert_data_test.py
@@ -65,17 +65,10 @@ def setup_key() -> tink.aead.KmsEnvelopeAead:
 
 
 def test_encrypt_and_insert_data(
-    capsys: pytest.CaptureFixture,
     pool: sqlalchemy.engine.Engine,
-    env_aead: tink.aead.KmsEnvelopeAead
+    env_aead: tink.aead.KmsEnvelopeAead,
 ) -> None:
-    encrypt_and_insert_data(
-        pool,
-        env_aead,
-        table_name,
-        "SPACES",
-        "hello@example.com")
-    captured = capsys.readouterr()
+    encrypt_and_insert_data(pool, env_aead, table_name, "SPACES", "hello@example.com")
 
     decrypted_emails = []
     with pool.connect() as conn:
@@ -89,5 +82,4 @@ def test_encrypt_and_insert_data(
             email = env_aead.decrypt(row[2], team.encode()).decode()
             decrypted_emails.append(email)
 
-    assert "Vote successfully cast for 'SPACES'" in captured.out
     assert "hello@example.com" in decrypted_emails
diff --git a/cloud-sql/postgres/client-side-encryption/snippets/query_and_decrypt_data.py b/cloud-sql/postgres/client-side-encryption/snippets/query_and_decrypt_data.py
@@ -24,6 +24,9 @@
 
 
 def main() -> None:
+    """
+    Connects to the database, inserts encrypted data and retrieves encrypted data.
+    """
     db_user = os.environ["DB_USER"]  # e.g. "root", "postgres"
     db_pass = os.environ["DB_PASS"]  # e.g. "mysupersecretpassword"
     db_name = os.environ["DB_NAME"]  # e.g. "votes_db"
@@ -67,7 +70,10 @@ def query_and_decrypt_data(
     db: sqlalchemy.engine.base.Engine,
     env_aead: tink.aead.KmsEnvelopeAead,
     table_name: str,
-) -> None:
+) -> list[tuple[str]]:
+    """
+    Retrieves data from the database and decrypts it using the KmsEnvelopeAead object.
+    """
     with db.connect() as conn:
         # Execute the query and fetch all results
         recent_votes = conn.execute(
@@ -76,6 +82,7 @@ def query_and_decrypt_data(
         ).fetchall()
 
         print("Team\tEmail\tTime Cast")
+        output = []
 
         for row in recent_votes:
             team = row[0]
@@ -93,6 +100,8 @@ def query_and_decrypt_data(
 
             # Print recent votes
             print(f"{team}\t{email}\t{time_cast}")
+            output.append((team, email, time_cast))
+    return output
 
 
 # [END cloud_sql_postgres_cse_query]
diff --git a/cloud-sql/postgres/client-side-encryption/snippets/query_and_decrypt_data_test.py b/cloud-sql/postgres/client-side-encryption/snippets/query_and_decrypt_data_test.py
@@ -68,19 +68,15 @@ def setup_key() -> tink.aead.KmsEnvelopeAead:
 def test_query_and_decrypt_data(
     capsys: pytest.CaptureFixture,
     pool: sqlalchemy.engine.Engine,
-    env_aead: tink.aead.KmsEnvelopeAead
+    env_aead: tink.aead.KmsEnvelopeAead,
 ) -> None:
-
     # Insert data into table before testing
-    encrypt_and_insert_data(
-        pool,
-        env_aead,
-        table_name,
-        "SPACES",
-        "hello@example.com")
-
-    query_and_decrypt_data(pool, env_aead, table_name)
-
-    captured = capsys.readouterr()
-    assert "Team\tEmail\tTime Cast" in captured.out
-    assert "hello@example.com" in captured.out
+    encrypt_and_insert_data(pool, env_aead, table_name, "SPACES", "hello@example.com")
+
+    output = query_and_decrypt_data(pool, env_aead, table_name)
+
+    for row in output:
+        if row[1] == "hello@example.com":
+            break
+    else:
+        pytest.fail("Failed to find vote in the decrypted data.")
