diff --git a/unit_tests/utilities/test_zaza_utilities_swift.py b/unit_tests/utilities/test_zaza_utilities_swift.py
index f0e330b..1fc2f5d 100644
--- a/unit_tests/utilities/test_zaza_utilities_swift.py
+++ b/unit_tests/utilities/test_zaza_utilities_swift.py
@@ -93,7 +93,7 @@ class TestSwiftUtils(ut_utils.BaseTestCase):
         self.patch_object(zaza.model, 'get_units')
         self.patch_object(zaza.model, 'get_unit_public_address')
 
-        def _get_unit_public_address(u):
+        def _get_unit_public_address(u, model_name=None):
             return u.public_address
 
         self.get_unit_public_address.side_effect = _get_unit_public_address
diff --git a/zaza/openstack/utilities/swift.py b/zaza/openstack/utilities/swift.py
index b7db584..cea0a33 100644
--- a/zaza/openstack/utilities/swift.py
+++ b/zaza/openstack/utilities/swift.py
@@ -207,7 +207,10 @@ def get_swift_storage_topology(model_name=None):
             region = app_config['storage-region']['value']
             zone = app_config['zone']['value']
             for unit in zaza.model.get_units(app_name, model_name=model_name):
-                topology[zaza.model.get_unit_public_address(unit)] = {
+                unit_ip = zaza.model.get_unit_public_address(
+                    unit,
+                    model_name=model_name)
+                topology[unit_ip] = {
                     'app_name': app_name,
                     'unit': unit,
                     'region': region,