Add juju max-frame-size option

docs/reference.md

@@ -79,6 +79,12 @@ Path to the directory where the `juju-crashdump` output will be stored. The defa
 the current working directory.
 
 
+### `--juju-max-frame-size`
+
+Maximum frame size to use when connecting to a juju model. The default is None and
+the value must be between 0 and MAX_FRAME_SIZE
+
+
 ## Fixtures
 
 ### `ops_test`

pytest_operator/plugin.py

@@ -22,6 +22,8 @@
 from string import ascii_lowercase, digits, hexdigits
 from timeit import default_timer as timer
 from typing import (
+    Any,
+    Dict,
     Generator,
     Iterable,
     List,
@@ -45,7 +47,7 @@
 import yaml
 from _pytest.config import Config
 from _pytest.config.argparsing import Parser
-from juju.client import client
+from juju.client import client, connection
 from juju.client.jujudata import FileJujuData
 from juju.errors import JujuError
 from juju.exceptions import DeadEntityException
@@ -143,6 +145,13 @@ def pytest_addoption(parser: Parser):
         "* ignored if `--model` supplied"
         "* if the specified file doesn't exist, an error will be raised.",
     )
+    parser.addoption(
+        "--juju-max-frame-size",
+        action="store",
+        default=None,
+        help="Set the maximum frame size for websocket communication with Juju.",
+        type=int,
+    )
 
 
 def pytest_load_initial_conftests(parser: Parser, args: List[str]) -> None:
@@ -414,6 +423,20 @@ class ModelInUseError(Exception):
 Timeout = TypeVar("Timeout", float, int)
 
 
+def _connect_kwds(request) -> Dict[str, Any]:
+    """Create a dict of keyword arguments for connecting to a model."""
+    kwds = {}
+    if val := request.config.option.juju_max_frame_size:
+        if 0 < val <= connection.Connection.MAX_FRAME_SIZE:
+            kwds["max_frame_size"] = val
+        else:
+            raise ValueError(
+                f"max-frame-size must be positive int and less than or equal to "
+                f"{connection.Connection.MAX_FRAME_SIZE}"
+            )
+    return kwds
+
+
 @dataclasses.dataclass
 class ModelState:
     model: Model
@@ -510,6 +533,7 @@ def __init__(self, request, tmp_path_factory):
         self._init_model_name: Optional[str] = request.config.option.model
         self._init_keep_model: bool = request.config.option.keep_models
         self._init_destroy_storage: bool = request.config.option.destroy_storage
+        self._juju_connect_kwds: Dict[str, Any] = _connect_kwds(request)
 
         # These may be modified by _setup_model
         self.controller_name = request.config.option.controller
@@ -753,7 +777,7 @@ async def _model_exists(self, model_name: str) -> bool:
 
     @staticmethod
     async def _connect_to_model(
-        controller_name, model_name, keep=True, destroy_storage=False
+        controller_name, model_name, keep=True, destroy_storage=False, **connect_kwargs
     ):
         """
         Makes a reference to an existing model used by the test framework
@@ -766,13 +790,13 @@ async def _connect_to_model(
         log.info(
             "Connecting to existing model %s on unspecified cloud", state.full_name
         )
-        await model.connect(state.full_name)
+        await model.connect(state.full_name, **connect_kwargs)
         state.config = await model.get_config()
         return state
 
     @staticmethod
     def read_model_config(
-        config_path_or_obj: Union[dict, str, os.PathLike, None]
+        config_path_or_obj: Union[dict, str, os.PathLike, None],
     ) -> Optional[dict]:
         if isinstance(config_path_or_obj, dict):
             return config_path_or_obj
@@ -796,7 +820,9 @@ async def _setup_model(self):
         assert self.controller_name, "No controller selected for ops_test"
         if not self._controller:
             self._controller = Controller()
-            await self._controller.connect(self.controller_name)
+            await self._controller.connect(
+                self.controller_name, **self._juju_connect_kwds
+            )
 
         await self.track_model(
             alias,
@@ -883,7 +909,10 @@ async def track_model(
                     "Cannot use_existing model if model_name is unspecified"
                 )
             model_state = await self._connect_to_model(
-                self.controller_name, model_name, keep_val
+                self.controller_name,
+                model_name,
+                keep_val,
+                **self._juju_connect_kwds,
             )
         else:
             cloud_name = cloud_name or self.cloud_name
@@ -1371,7 +1400,7 @@ async def async_render_bundles(self, *bundles: BundleOpt, **context) -> List[Pat
                 bundle_zip = ZipPath(filepath, "bundle.yaml")
                 content = bundle_zip.read_text()
             else:
-                raise TypeError("bundle {} isn't a known Type".format(type(bundle)))
+                raise TypeError(f"bundle {type(bundle)} isn't a known Type")
             to_render.append(content)
         return self.render_bundles(*to_render, **context)
 

ruff.toml

@@ -0,0 +1,51 @@
+# Support Python 3.8.
+target-version = "py38"
+# Set the maximum line length to 88
+line-length = 88
+indent-width = 4
+
+[lint]
+# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`)  codes by default.
+# Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or
+# McCabe complexity (`C901`) by default.
+select = ["E4", "E7", "E9", "F", "B", "Q"]
+ignore = []
+
+# Allow fix for all enabled rules (when `--fix`) is provided.
+fixable = ["ALL"]
+unfixable = []
+
+# Allow unused variables when underscore-prefixed.
+# wokeignore:rule=dummy
+dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+
+[lint.pydocstyle]
+convention = "google"
+
+
+[format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
+
+# Enable auto-formatting of code examples in docstrings. Markdown,
+# reStructuredText code/literal blocks and doctests are all supported.
+#
+# This is currently disabled by default, but it is planned for this
+# to be opt-out in the future.
+docstring-code-format = false
+
+# Set the line length limit used when formatting code snippets in
+# docstrings.
+#
+# This only has an effect when the `docstring-code-format` setting is
+# enabled.
+docstring-code-line-length = "dynamic"

setup.py

@@ -9,7 +9,7 @@
     author_email="[email protected]",
     description="Fixtures for Operators",
     long_description=Path("README.md").read_text(),
-    long_description_content_type='text/markdown',
+    long_description_content_type="text/markdown",
     classifiers=[
         "Framework :: Pytest",
         "Programming Language :: Python",
@@ -25,9 +25,9 @@
     keywords=["pytest", "py.test", "operators", "yaml", "asyncio"],
     name="pytest-operator",
     packages=find_packages(include=["pytest_operator"]),
-    package_data={'pytest_operator': ['py.typed']},
+    package_data={"pytest_operator": ["py.typed"]},
     url="https://github.com/charmed-kubernetes/pytest-operator",
-    version="0.38.0",
+    version="0.39.0",
     zip_safe=True,
     install_requires=[
         "ipdb",

tests/unit/test_pytest_operator.py

@@ -60,13 +60,13 @@ async def test_without_tox(request, ops_test):
     result.assert_outcomes(passed=1)
 
 
-async def test_destructive_mode(monkeypatch, tmp_path_factory):
+async def test_destructive_mode(setup_request, monkeypatch, tmp_path_factory):
     patch = monkeypatch.setattr
     patch(plugin.os, "getgroups", mock_getgroups := Mock(return_value=[]))
     patch(plugin.grp, "getgrall", mock_getgrall := Mock(return_value=[]))
     patch(plugin.grp, "getgrgid", Mock(return_value=Mock(gr_name="lxd")))
     patch(plugin.OpsTest, "run", mock_run := AsyncMock(return_value=(1, "", "")))
-    ops_test = plugin.OpsTest(Mock(**{"module.__name__": "test"}), tmp_path_factory)
+    ops_test = plugin.OpsTest(setup_request, tmp_path_factory)
 
     ops_test.destructive_mode = True
     try:
@@ -225,8 +225,8 @@ async def resource_charm(request, tmp_path_factory):
     yield dst_path
 
 
-async def test_plugin_build_resources(tmp_path_factory):
-    ops_test = plugin.OpsTest(Mock(**{"module.__name__": "test"}), tmp_path_factory)
+async def test_plugin_build_resources(setup_request, tmp_path_factory):
+    ops_test = plugin.OpsTest(setup_request, tmp_path_factory)
     ops_test.jujudata = Mock()
     ops_test.jujudata.path = ""
     dst_dir = ops_test.tmp_path / "resources"
@@ -257,8 +257,8 @@ async def test_plugin_build_resources(tmp_path_factory):
     assert resources and resources == expected_resources
 
 
-async def test_plugin_get_resources(tmp_path_factory, resource_charm):
-    ops_test = plugin.OpsTest(Mock(**{"module.__name__": "test"}), tmp_path_factory)
+async def test_plugin_get_resources(setup_request, tmp_path_factory, resource_charm):
+    ops_test = plugin.OpsTest(setup_request, tmp_path_factory)
     resources = ops_test.arch_specific_resources(resource_charm)
     assert resources.keys() == {"resource-file-arm64", "resource-file"}
     assert resources["resource-file-arm64"].arch == "arm64"
@@ -269,8 +269,8 @@ async def test_plugin_get_resources(tmp_path_factory, resource_charm):
     "pytest_operator.plugin.CharmStore._charm_id",
     new=Mock(return_value="resourced-charm-1"),
 )
-async def test_plugin_fetch_resources(tmp_path_factory, resource_charm):
-    ops_test = plugin.OpsTest(Mock(**{"module.__name__": "test"}), tmp_path_factory)
+async def test_plugin_fetch_resources(setup_request, tmp_path_factory, resource_charm):
+    ops_test = plugin.OpsTest(setup_request, tmp_path_factory)
     ops_test.jujudata = Mock()
     ops_test.jujudata.path = ""
     arch_resources = ops_test.arch_specific_resources(resource_charm)
@@ -295,8 +295,8 @@ def dl_rsc(resource, dest_path):
     assert downloaded == expected_downloads
 
 
-async def test_async_render_bundles(tmp_path_factory):
-    ops_test = plugin.OpsTest(Mock(**{"module.__name__": "test"}), tmp_path_factory)
+async def test_async_render_bundles(setup_request, tmp_path_factory):
+    ops_test = plugin.OpsTest(setup_request, tmp_path_factory)
     ops_test.jujudata = Mock()
     ops_test.jujudata.path = ""
 
@@ -358,6 +358,7 @@ async def test_async_render_bundles(tmp_path_factory):
     ],
 )
 async def test_crash_dump_mode(
+    setup_request,
     crash_dump,
     no_crash_dump,
     n_testsfailed,
@@ -369,12 +370,11 @@ async def test_crash_dump_mode(
     """Test running juju-crashdump in OpsTest.cleanup."""
     patch = monkeypatch.setattr
     patch(plugin.OpsTest, "run", mock_run := AsyncMock(return_value=(0, "", "")))
-    mock_request = Mock(**{"module.__name__": "test"})
-    mock_request.config.option.crash_dump = crash_dump
-    mock_request.config.option.no_crash_dump = no_crash_dump
-    mock_request.config.option.crash_dump_args = "-c --bug=1234567"
-    mock_request.config.option.keep_models = False
-    ops_test = plugin.OpsTest(mock_request, tmp_path_factory)
+    setup_request.config.option.crash_dump = crash_dump
+    setup_request.config.option.no_crash_dump = no_crash_dump
+    setup_request.config.option.crash_dump_args = "-c --bug=1234567"
+    setup_request.config.option.keep_models = False
+    ops_test = plugin.OpsTest(setup_request, tmp_path_factory)
     model = MagicMock()
     model.machines.values.return_value = []
     model.disconnect = AsyncMock()
@@ -389,7 +389,7 @@ async def test_crash_dump_mode(
     ops_test.log_model = AsyncMock()
     ops_test._controller = AsyncMock()
 
-    mock_request.session.testsfailed = n_testsfailed
+    setup_request.session.testsfailed = n_testsfailed
 
     await ops_test._cleanup_model()
 
@@ -407,18 +407,17 @@ async def test_crash_dump_mode(
         mock_run.assert_not_called()
 
 
-def test_crash_dump_mode_invalid_input(monkeypatch, tmp_path_factory):
+def test_crash_dump_mode_invalid_input(setup_request, monkeypatch, tmp_path_factory):
     """Test running juju-crashdump in OpsTest.cleanup."""
     patch = monkeypatch.setattr
     patch(plugin.OpsTest, "run", AsyncMock(return_value=(0, "", "")))
-    mock_request = Mock(**{"module.__name__": "test"})
-    mock_request.config.option.crash_dump = "not-a-real-option"
-    mock_request.config.option.crash_dump_args = ""
-    mock_request.config.option.no_crash_dump = False
-    mock_request.config.option.keep_models = False
+    setup_request.config.option.crash_dump = "not-a-real-option"
+    setup_request.config.option.crash_dump_args = ""
+    setup_request.config.option.no_crash_dump = False
+    setup_request.config.option.keep_models = False
 
     with pytest.raises(ValueError):
-        plugin.OpsTest(mock_request, tmp_path_factory)
+        plugin.OpsTest(setup_request, tmp_path_factory)
 
 
 async def test_create_crash_dump(monkeypatch, tmp_path_factory):
@@ -432,6 +431,7 @@ async def mock_run(*cmd):
     patch(plugin.OpsTest, "run", mock_run)
     mock_request = Mock(**{"module.__name__": "test"})
     mock_request.config.option.crash_dump_args = ""
+    mock_request.config.option.juju_max_frame_size = None
     patch(plugin, "log", mock_log := MagicMock())
     ops_test = plugin.OpsTest(mock_request, tmp_path_factory)
     await ops_test.create_crash_dump()
@@ -504,19 +504,28 @@ def setup_request(request, mock_juju):
     mock_request.config.option.model_config = None
     mock_request.config.option.keep_models = False
     mock_request.config.option.destroy_storage = False
+    mock_request.config.option.juju_max_frame_size = None
     yield mock_request
 
 
+@pytest.mark.parametrize("max_frame_size", [None, 2**16])
 async def test_fixture_set_up_existing_model(
-    mock_juju, setup_request, tmp_path_factory
+    mock_juju, setup_request, tmp_path_factory, max_frame_size
 ):
     setup_request.config.option.model = "this-model"
+    setup_request.config.option.juju_max_frame_size = max_frame_size
+    expected_kwargs = {}
+    if max_frame_size:
+        expected_kwargs["max_frame_size"] = max_frame_size
+
     ops_test = plugin.OpsTest(setup_request, tmp_path_factory)
     assert ops_test.model is None
 
     mock_juju.controller.list_models = AsyncMock(return_value=["this-model"])
     await ops_test._setup_model()
-    mock_juju.model.connect.assert_called_with("this-controller:this-model")
+    mock_juju.model.connect.assert_called_with(
+        "this-controller:this-model", **expected_kwargs
+    )
     assert ops_test.model == mock_juju.model
     assert ops_test.model_full_name == "this-controller:this-model"
     assert ops_test.cloud_name is None
@@ -528,6 +537,16 @@ async def test_fixture_set_up_existing_model(
     assert len(ops_test.models) == 1
 
 
+@pytest.mark.parametrize("max_frame_size", [-1, 2**32])
+async def test_fixture_invalid_max_frame_size(
+    setup_request, tmp_path_factory, max_frame_size
+):
+    setup_request.config.option.model = "this-model"
+    setup_request.config.option.juju_max_frame_size = max_frame_size
+    with pytest.raises(ValueError):
+        plugin.OpsTest(setup_request, tmp_path_factory)
+
+
 @patch("pytest_operator.plugin.OpsTest.forget_model")
 @patch("pytest_operator.plugin.OpsTest.run")
 async def test_fixture_cleanup_multi_model(