From 0d08092227654fbc396e3016a137b0d2d5afc886 Mon Sep 17 00:00:00 2001
From: Adam Dyess <adam.dyess@canonical.com>
Date: Fri, 26 Jul 2024 12:49:18 -0500
Subject: [PATCH] unify parse_storage_constraint or client.Constraints during
 deploy

---
 juju/model.py                   | 11 +++++++++--
 tests/integration/test_model.py | 15 +++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/juju/model.py b/juju/model.py
index a51944de3..66211cc11 100644
--- a/juju/model.py
+++ b/juju/model.py
@@ -2105,6 +2105,13 @@ async def _deploy(self, charm_url, application, series, config,
 
         app_facade = client.ApplicationFacade.from_connection(self.connection())
 
+        # Prepare all storage constraints
+        storage = storage or dict()
+        storage = {
+            k: v if isinstance(v, client.Constraints) else parse_storage_constraint(v)
+            for k, v in storage.items()
+        }
+
         if server_side_deploy:
             # Call DeployFromRepository
             app = client.DeployFromRepositoryArg(
@@ -2116,7 +2123,7 @@ async def _deploy(self, charm_url, application, series, config,
                 devices=devices,
                 dryrun=False,
                 placement=placement,
-                storage={k: parse_storage_constraint(v) for k, v in (storage or dict()).items()},
+                storage=storage,
                 trust=trust,
                 base=charm_origin.base,
                 channel=channel,
@@ -2151,7 +2158,7 @@ async def _deploy(self, charm_url, application, series, config,
                 endpoint_bindings=endpoint_bindings,
                 num_units=num_units,
                 resources=resources,
-                storage={k: parse_storage_constraint(v) for k, v in (storage or dict()).items()},
+                storage=storage,
                 placement=placement,
                 devices=devices,
                 attach_storage=attach_storage,
diff --git a/tests/integration/test_model.py b/tests/integration/test_model.py
index 8d213553a..c488b7502 100644
--- a/tests/integration/test_model.py
+++ b/tests/integration/test_model.py
@@ -1197,6 +1197,21 @@ async def test_model_cache_update():
         await controller.destroy_models(model_name)
 
 
+@base.bootstrapped
+async def test_deploy_with_storage():
+    async with base.CleanModel() as model:
+        await model.deploy(
+            'postgresql',
+            storage={"pgdata": {"size": 1024, "count": 1}},
+        )
+        await model.wait_for_idle(status="active")
+        storages = await model.list_storage()
+        await model.list_storage(filesystem=True)
+        await model.list_storage(volume=True)
+
+        assert any([tag.storage("pgdata") in s['storage-tag'] for s in storages])
+
+
 @base.bootstrapped
 async def test_add_storage():
     pytest.skip('skip in favour of test_add_and_list_storage')