From 5db5f1cd96aebc36193f525724626598e52fcb96 Mon Sep 17 00:00:00 2001
From: Daniel Goldman <danielgoldman4@gmail.com>
Date: Fri, 5 Jul 2024 17:38:13 -0400
Subject: [PATCH] add more test

---
 llamazure/azrest/azrest_test.py    | 55 +++++++++++++++++++++++++++++-
 llamazure/rbac/integration_test.py | 11 +++---
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/llamazure/azrest/azrest_test.py b/llamazure/azrest/azrest_test.py
index 3d73ccc..9000b57 100644
--- a/llamazure/azrest/azrest_test.py
+++ b/llamazure/azrest/azrest_test.py
@@ -3,9 +3,10 @@
 
 # pylint: disable=redefined-outer-name
 import pytest
+from pydantic import BaseModel, ValidationError
 
 from llamazure.azrest.azrest import AzRest
-from llamazure.azrest.models import AzList, AzureError, BatchReq, Req
+from llamazure.azrest.models import AzList, AzureError, BatchReq, Req, cast_as, ensure
 
 
 @pytest.fixture
@@ -89,3 +90,55 @@ def test_custom_names(self, azr, it_info):
 			assert isinstance(res, AzList)
 
 		assert len(batch_res["test-req-0"].value) > 0
+
+
+# Example Pydantic models
+class Foo(BaseModel):
+	id: int
+	name: str
+
+
+class FooUpdateParameters(BaseModel):
+	name: str
+
+
+class TestCastAs:
+	def test_cast_foo_to_fooupdateparameters(self):
+		foo = Foo(id=1, name="test")
+		foo_update = cast_as(foo, FooUpdateParameters)
+		assert foo_update.name == foo.name
+		assert not hasattr(foo_update, "id")
+
+	def test_cast_fooupdateparameters_to_foo(self):
+		foo_update = FooUpdateParameters(name="updated_name")
+		with pytest.raises(ValidationError):
+			cast_as(foo_update, Foo)  # This should raise an error since 'id' is missing
+
+	def test_cast_with_additional_fields(self):
+		class Bar(BaseModel):
+			name: str
+			extra_field: int
+
+		foo = Foo(id=1, name="test")
+		with pytest.raises(ValueError):
+			cast_as(foo, Bar)  # This should raise an error since 'extra_field' is missing
+
+	def test_cast_with_missing_fields(self):
+		class FooPartial(BaseModel):
+			id: int
+
+		foo = Foo(id=1, name="test")
+		foo_partial = cast_as(foo, FooPartial)
+		assert foo_partial.id == foo.id
+		assert not hasattr(foo_partial, "name")
+
+
+class TestEnsure:
+	def test_ensure_not_none(self):
+		assert ensure(5) == 5
+		assert ensure("test") == "test"
+		assert ensure([1, 2, 3]) == [1, 2, 3]
+
+	def test_ensure_none_raises(self):
+		with pytest.raises(TypeError, match="value was None"):
+			ensure(None)
diff --git a/llamazure/rbac/integration_test.py b/llamazure/rbac/integration_test.py
index 54c243e..1cfb8d4 100644
--- a/llamazure/rbac/integration_test.py
+++ b/llamazure/rbac/integration_test.py
@@ -76,11 +76,14 @@ def assert_role_assigned():
 		# explicitly make a `put` that already exists
 		retry(assert_role_assigned, AzureError)
 
-		l.info("deleting RoleAssignment")
-		ras.DeleteById(asn.rid)
+		def cleanup():
+			l.info("deleting RoleAssignment")
+			ras.DeleteById(asn.rid)
 
-		l.info("cleanup role")
-		retry(lambda: rds.delete_by_name(role.properties.roleName), AzureError)
+			l.info("cleanup role")
+			rds.delete_by_name(role.properties.roleName)
+
+		retry(lambda: cleanup, AzureError)
 
 	@pytest.mark.integration
 	@pytest.mark.admin