Update databag model dump return value

src/cosl/coordinated_workers/interface.py

@@ -94,26 +94,27 @@ def load(cls, databag: _RawDatabag):
             log.debug(msg, exc_info=True)
             raise DataValidationError(msg) from e
 
-    def dump(self, databag: Optional[_RawDatabag] = None, clear: bool = True) -> None:
+    def dump(self, databag: Optional[_RawDatabag] = None, clear: bool = True) -> _RawDatabag:
         """Write the contents of this model to Juju databag.
 
         :param databag: the databag to write the data to.
         :param clear: ensure the databag is cleared before writing it.
         """
-        if clear and databag:
-            databag.clear()
+        _databag: _RawDatabag = {} if databag is None else databag
+
+        if clear:
+            _databag.clear()
 
-        if databag is None:
-            databag = {}
         if nest_under := self.model_config.get("_NEST_UNDER"):
-            databag[nest_under] = self.model_dump_json(  # type: ignore
+            _databag[nest_under] = self.model_dump_json(  # type: ignore
                 by_alias=True,
                 # skip keys whose values are default
                 exclude_defaults=True,
             )
 
         dct = self.model_dump(mode="json", by_alias=True, exclude_defaults=True)  # type: ignore
-        databag.update({k: json.dumps(v) for k, v in dct.items()})
+        _databag.update({k: json.dumps(v) for k, v in dct.items()})
+        return _databag
 
 
 # =============

tests/test_coordinated_workers/test_interface.py

@@ -0,0 +1,19 @@
+from cosl.coordinated_workers.interface import DatabagModel
+
+
+def test_databag_dump_update():
+    class MyModel(DatabagModel):
+        foo: int
+        bar: str = "barian"
+
+    db = {}
+    assert MyModel(foo=1, bar="soz").dump(db)
+    assert db == {"foo": "1", "bar": '"soz"'}
+
+
+def test_databag_dumps():
+    class MyModel(DatabagModel):
+        foo: int
+        bar: str = "barian"
+
+    assert MyModel(foo=1, bar="bearian").dump() == {"foo": "1", "bar": '"bearian"'}