update tests to work for new vault object types

Author: debussyman

tests/test_vault.py

@@ -2,8 +2,8 @@
 from tests.utils.fixtures.mock_vault_object import (
     MockVaultObject,
     MockObjectVersion,
-    MockDataKey,
-    MockDataKeyPair,
+    MockObjectDigest,
+    MockObjectMetadata,
 )
 from tests.utils.list_resource import list_response_of
 from tests.utils.syncify import syncify
@@ -23,6 +23,14 @@ def mock_vault_object(self):
             "vault_01234567890abcdef", "test-secret", "secret-value"
         ).dict()
 
+    @pytest.fixture
+    def mock_object_digest(self):
+        return MockObjectDigest("vault_01234567890abcdef", "test-secret").dict()
+
+    @pytest.fixture
+    def mock_object_metadata(self):
+        return MockObjectMetadata("vault_01234567890abcdef").dict()
+
     @pytest.fixture
     def mock_vault_object_no_value(self):
         mock_obj = MockVaultObject("vault_01234567890abcdef", "test-secret")
@@ -32,8 +40,7 @@ def mock_vault_object_no_value(self):
     @pytest.fixture
     def mock_vault_objects_list(self):
         vault_objects = [
-            MockVaultObject(f"vault_{i}", f"secret-{i}", f"value-{i}").dict()
-            for i in range(5)
+            MockObjectDigest(f"vault_{i}", f"secret-{i}").dict() for i in range(5)
         ]
         return {
             "data": vault_objects,
@@ -44,8 +51,7 @@ def mock_vault_objects_list(self):
     @pytest.fixture
     def mock_vault_objects_multiple_pages(self):
         vault_objects = [
-            MockVaultObject(f"vault_{i}", f"secret-{i}", f"value-{i}").dict()
-            for i in range(25)
+            MockObjectDigest(f"vault_{i}", f"secret-{i}").dict() for i in range(25)
         ]
         return list_response_of(data=vault_objects)
 
@@ -176,13 +182,13 @@ def test_list_object_versions_empty_data(
         assert len(versions) == 0
 
     def test_create_object_success(
-        self, mock_vault_object, capture_and_mock_http_client_request
+        self, mock_object_metadata, capture_and_mock_http_client_request
     ):
         request_kwargs = capture_and_mock_http_client_request(
-            self.http_client, mock_vault_object, 200
+            self.http_client, mock_object_metadata, 200
         )
 
-        vault_object = self.vault.create_object(
+        object_metadata = self.vault.create_object(
             name="test-secret",
             value="secret-value",
             key_context=KeyContext({"key": "test-key"}),
@@ -192,10 +198,8 @@ def test_create_object_success(
         assert request_kwargs["url"].endswith("/vault/v1/kv")
         assert request_kwargs["json"]["name"] == "test-secret"
         assert request_kwargs["json"]["value"] == "secret-value"
-        assert request_kwargs["json"]["context"] == KeyContext({"key": "test-key"})
-        assert vault_object.id == "vault_01234567890abcdef"
-        assert vault_object.name == "test-secret"
-        assert vault_object.value == "secret-value"
+        assert request_kwargs["json"]["key_context"] == KeyContext({"key": "test-key"})
+        assert object_metadata.id == "vault_01234567890abcdef"
 
     def test_create_object_missing_name(self):
         with pytest.raises(

tests/utils/fixtures/mock_vault_object.py

@@ -2,6 +2,7 @@
 
 from workos.types.vault import (
     VaultObject,
+    ObjectDigest,
     ObjectMetadata,
     ObjectUpdateBy,
     ObjectVersion,
@@ -37,38 +38,31 @@ def __init__(
         )
 
 
-class MockObjectVersion(ObjectVersion):
-    def __init__(self, id="version_01234567890abcdef", current_version=True):
+class MockObjectDigest(ObjectDigest):
+    def __init__(self, id="vault_01234567890abcdef", name="test-secret"):
         now = datetime.datetime.now().isoformat()
-        super().__init__(
-            id=id,
-            created_at=now,
-            current_version=current_version,
-        )
+        super().__init__(id=id, name=name, updated_at=now)
 
 
-class MockDataKey(DataKey):
-    def __init__(
-        self,
-        id="key_01234567890abcdef",
-        key="MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=",
-    ):
+class MockObjectMetadata(ObjectMetadata):
+    def __init__(self, id="vault_01234567890abcdef"):
+        now = datetime.datetime.now().isoformat()
         super().__init__(
+            context=KeyContext(key="test-key"),
+            environment_id="env_01234567890abcdef",
             id=id,
-            key=key,
+            key_id="key_01234567890abcdef",
+            updated_at=now,
+            updated_by=ObjectUpdateBy(id="user_01234567890abcdef", name="Test User"),
+            version_id="version_01234567890abcdef",
         )
 
 
-class MockDataKeyPair(DataKeyPair):
-    def __init__(
-        self, context=None, data_key=None, encrypted_keys="ZW5jcnlwdGVkX2tleXNfZGF0YQ=="
-    ):
-        if context is None:
-            context = VaultKeyContext({"key": "test-key"})
-        if data_key is None:
-            data_key = MockDataKey()
+class MockObjectVersion(ObjectVersion):
+    def __init__(self, id="version_01234567890abcdef", current_version=True):
+        now = datetime.datetime.now().isoformat()
         super().__init__(
-            context=context,
-            data_key=data_key,
-            encrypted_keys=encrypted_keys,
+            id=id,
+            created_at=now,
+            current_version=current_version,
         )

workos/types/list_resource.py

@@ -33,7 +33,7 @@
 from workos.types.organizations import Organization
 from workos.types.sso import ConnectionWithDomains
 from workos.types.user_management import Invitation, OrganizationMembership, User
-from workos.types.vault import VaultObject
+from workos.types.vault import ObjectDigest
 from workos.types.workos_model import WorkOSModel
 from workos.utils.request_helper import DEFAULT_LIST_RESPONSE_LIMIT
 
@@ -52,7 +52,7 @@
     AuthorizationResource,
     AuthorizationResourceType,
     User,
-    VaultObject,
+    ObjectDigest,
     Warrant,
     WarrantQueryResult,
 )

workos/vault.py

@@ -246,13 +246,14 @@ def list_objects(
         )
 
         # Ensure object field is present
-        if "object" not in response:
-            response["object"] = "list"
+        response_dict = dict(response)
+        if "object" not in response_dict:
+            response_dict["object"] = "list"
 
         return VaultObjectList(
             list_method=self.list_objects,
             list_args=list_params,
-            **ListPage[ObjectDigest](**response).model_dump(),
+            **ListPage[ObjectDigest](**response_dict).model_dump(),
         )
 
     def list_object_versions(