fix tests

Author: nidhiii-27

samples/snippets/encryption_test.py

@@ -22,6 +22,7 @@
 from google.cloud.storage import Blob
 import pytest
 
+from google.cloud.storage.bucket import EncryptionEnforcementConfig
 import storage_download_encrypted_file
 import storage_generate_encryption_key
 import storage_object_csek_to_cmek
@@ -88,11 +89,7 @@ def test_blob():
     except NotFound as e:
         # For the case that the rotation succeeded.
         print(f"Ignoring 404, detail: {e}")
-        blob = Blob(
-            blob_name,
-            bucket,
-            encryption_key=TEST_ENCRYPTION_KEY_2_DECODED
-        )
+        blob = Blob(blob_name, bucket, encryption_key=TEST_ENCRYPTION_KEY_2_DECODED)
         blob.delete()
 
 
@@ -152,27 +149,64 @@ def test_set_bucket_encryption_enforcement_config(enforcement_bucket):
     storage_client = storage.Client()
     bucket = storage_client.get_bucket(enforcement_bucket)
 
-    assert bucket.google_managed_encryption_enforcement_config.restriction_mode == "FullyRestricted"
-    assert bucket.customer_managed_encryption_enforcement_config.restriction_mode == "NotRestricted"
-    assert bucket.customer_supplied_encryption_enforcement_config.restriction_mode == "FullyRestricted"
+    assert (
+        bucket.encryption.google_managed_encryption_enforcement_config.restriction_mode
+        == "FullyRestricted"
+    )
+    assert (
+        bucket.encryption.customer_managed_encryption_enforcement_config.restriction_mode
+        == "NotRestricted"
+    )
+    assert (
+        bucket.encryption.customer_supplied_encryption_enforcement_config.restriction_mode
+        == "FullyRestricted"
+    )
 
 
-def test_get_bucket_encryption_enforcement_config(enforcement_bucket):
-    # This just exercises the get snippet. If it crashes, the test fails.
-    # The assertions on the state were done in the set test.
+def test_get_bucket_encryption_enforcement_config(enforcement_bucket, capsys):
     storage_get_bucket_encryption_enforcement_config.get_bucket_encryption_enforcement_config(
         enforcement_bucket
     )
 
+    out, _ = capsys.readouterr()
+    assert f"Encryption Enforcement Config for bucket {enforcement_bucket}" in out
+    assert (
+        "Customer-managed encryption enforcement config restriction mode: NotRestricted"
+        in out
+    )
+    assert (
+        "Customer-supplied encryption enforcement config restriction mode: FullyRestricted"
+        in out
+    )
+    assert (
+        "Google-managed encryption enforcement config restriction mode: FullyRestricted"
+        in out
+    )
+
 
 def test_update_encryption_enforcement_config(enforcement_bucket):
+    storage_client = storage.Client()
+
+    # Pre-condition: Ensure bucket is in a different state before update
+    bucket = storage_client.get_bucket(enforcement_bucket)
+    bucket.encryption.google_managed_encryption_enforcement_config.restriction_mode = (
+        "NotRestricted"
+    )
+    bucket.patch()
+
     storage_update_encryption_enforcement_config.update_encryption_enforcement_config(
         enforcement_bucket
     )
 
     storage_client = storage.Client()
     bucket = storage_client.get_bucket(enforcement_bucket)
 
-    assert bucket.google_managed_encryption_enforcement_config.restriction_mode == "FullyRestricted"
-    assert bucket.customer_managed_encryption_enforcement_config.restriction_mode == "NotRestricted"
-    assert bucket.customer_supplied_encryption_enforcement_config is None
+    assert (
+        bucket.encryption.google_managed_encryption_enforcement_config.restriction_mode
+        == "FullyRestricted"
+    )
+    assert (
+        bucket.encryption.customer_managed_encryption_enforcement_config.restriction_mode
+        == "NotRestricted"
+    )
+    assert bucket.encryption.customer_supplied_encryption_enforcement_config is None

samples/snippets/storage_get_bucket_encryption_enforcement_config.py

@@ -26,13 +26,21 @@ def get_bucket_encryption_enforcement_config(bucket_name):
 
     print(f"Encryption Enforcement Config for bucket {bucket.name}:")
 
-    cmek_config = bucket.customer_managed_encryption_enforcement_config
-    csek_config = bucket.customer_supplied_encryption_enforcement_config
-    gmek_config = bucket.google_managed_encryption_enforcement_config
+    cmek_config = bucket.encryption.customer_managed_encryption_enforcement_config
+    csek_config = bucket.encryption.customer_supplied_encryption_enforcement_config
+    gmek_config = bucket.encryption.google_managed_encryption_enforcement_config
+
+    print(
+        f"Customer-managed encryption enforcement config restriction mode: {cmek_config.restriction_mode if cmek_config else None}"
+    )
+    print(
+        f"Customer-supplied encryption enforcement config restriction mode: {csek_config.restriction_mode if csek_config else None}"
+    )
+    print(
+        f"Google-managed encryption enforcement config restriction mode: {gmek_config.restriction_mode if gmek_config else None}"
+    )
+
 
-    print(f"Customer-managed encryption enforcement config restriction mode: {cmek_config.restriction_mode if cmek_config else None}")
-    print(f"Customer-supplied encryption enforcement config restriction mode: {csek_config.restriction_mode if csek_config else None}")
-    print(f"Google-managed encryption enforcement config restriction mode: {gmek_config.restriction_mode if gmek_config else None}")
 # [END storage_get_bucket_encryption_enforcement_config]
 
 

samples/snippets/storage_set_bucket_encryption_enforcement_config.py

@@ -27,25 +27,27 @@ def set_bucket_encryption_enforcement_config(bucket_name):
 
     # Setting restriction_mode to "FullyRestricted" for Google-managed encryption (GMEK)
     # means objects cannot be created using the default Google-managed keys.
-    bucket.google_managed_encryption_enforcement_config = EncryptionEnforcementConfig(
-        restriction_mode="FullyRestricted"
+    bucket.encryption.google_managed_encryption_enforcement_config = (
+        EncryptionEnforcementConfig(restriction_mode="FullyRestricted")
     )
 
     # Setting restriction_mode to "NotRestricted" for Customer-managed encryption (CMEK)
     # ensures that objects ARE permitted to be created using Cloud KMS keys.
-    bucket.customer_managed_encryption_enforcement_config = EncryptionEnforcementConfig(
-        restriction_mode="NotRestricted"
+    bucket.encryption.customer_managed_encryption_enforcement_config = (
+        EncryptionEnforcementConfig(restriction_mode="NotRestricted")
     )
 
     # Setting restriction_mode to "FullyRestricted" for Customer-supplied encryption (CSEK)
     # prevents objects from being created using raw, client-side provided keys.
-    bucket.customer_supplied_encryption_enforcement_config = EncryptionEnforcementConfig(
-        restriction_mode="FullyRestricted"
+    bucket.encryption.customer_supplied_encryption_enforcement_config = (
+        EncryptionEnforcementConfig(restriction_mode="FullyRestricted")
     )
 
     bucket.create()
 
     print(f"Created bucket {bucket.name} with Encryption Enforcement Config.")
+
+
 # [END storage_set_bucket_encryption_enforcement_config]
 
 

samples/snippets/storage_update_encryption_enforcement_config.py

@@ -19,23 +19,30 @@
 
 def update_encryption_enforcement_config(bucket_name):
     """Updates the encryption enforcement policy for a bucket."""
-    # The ID of your GCS bucket
+    # The ID of your GCS bucket with CMEK restricted
     # bucket_name = "your-unique-bucket-name"
 
     storage_client = storage.Client()
     bucket = storage_client.get_bucket(bucket_name)
 
-    # 1. Update a specific type (e.g., change GMEK to FullyRestricted)
-    bucket.google_managed_encryption_enforcement_config = EncryptionEnforcementConfig(restriction_mode="FullyRestricted")
-    bucket.customer_managed_encryption_enforcement_config = EncryptionEnforcementConfig(restriction_mode="NotRestricted")
+    # Update a specific type (e.g., change GMEK to FullyRestricted)
+    bucket.encryption.google_managed_encryption_enforcement_config = (
+        EncryptionEnforcementConfig(restriction_mode="FullyRestricted")
+    )
 
-    # 2. Remove a specific type (e.g., remove CSEK enforcement)
-    bucket.customer_supplied_encryption_enforcement_config = None
+    # Update another type (e.g., change CMEK to NotRestricted)
+    bucket.encryption.customer_managed_encryption_enforcement_config = (
+        EncryptionEnforcementConfig(restriction_mode="NotRestricted")
+    )
 
     bucket.patch()
 
     print(f"Encryption enforcement policy updated for bucket {bucket.name}.")
-    print("GMEK is now fully restricted, CMEK is now not restricted, and CSEK enforcement has been removed.")
+    print(
+        "GMEK is now fully restricted, CMEK is now not restricted, and CSEK enforcement has been removed."
+    )
+
+
 # [END storage_update_encryption_enforcement_config]