Fix using ext functions in MQTT publish (#851)

example/mqtt/test_mqtt.tavern.yaml

@@ -740,3 +740,28 @@ stages:
       payload: "there"
       timeout: 5
       qos: 1
+
+---
+
+test_name: Update an MQTT publish from an ext function
+
+includes:
+  - !include common.yaml
+
+paho-mqtt: *mqtt_spec
+
+stages:
+  - *setup_device_for_test
+
+  - name: step 1 - ping/pong
+    mqtt_publish:
+      topic: /device/{random_device_id}/echo
+      json:
+        $ext:
+          function: testing_utils:return_hello
+    mqtt_response:
+        topic: /device/{random_device_id}/echo/response
+        timeout: 3
+        qos: 1
+        json:
+          hello: there

example/mqtt/testing_utils.py

@@ -3,5 +3,5 @@ def message_says_hello(msg):
     assert msg.payload.get("message") == "hello world"
 
 
-def return_hello(_):
+def return_hello(_=None):
     return {"hello": "there"}

tavern/_core/exceptions.py

@@ -114,7 +114,7 @@ class InvalidFormattedJsonError(TavernException):
     """Tried to use the magic json format tag in an invalid way"""
 
 
-class InvalidExtBlockException(TavernException):
+class MisplacedExtBlockException(TavernException):
     """Tried to use the '$ext' block in a place it is no longer valid to use it"""
 
     def __init__(self, block) -> None:

tavern/_core/extfunctions.py

@@ -123,6 +123,11 @@ def inner():
 
 
 def _get_ext_values(ext: Mapping):
+    if not isinstance(ext, Mapping):
+        raise exceptions.InvalidExtFunctionError(
+            "ext block should be a dict, but it was a {}".format(type(ext))
+        )
+
     args = ext.get("extra_args") or ()
     kwargs = ext.get("extra_kwargs") or {}
     try:
@@ -145,14 +150,23 @@ def update_from_ext(request_args: dict, keys_to_check: List[str]) -> None:
     """
 
     new_args = {}
+    logger = _getlogger()
 
     for key in keys_to_check:
         try:
-            func = get_wrapped_create_function(request_args[key].pop("$ext"))
-        except (KeyError, TypeError, AttributeError):
-            pass
-        else:
-            new_args[key] = func()
+            block = request_args[key]
+        except KeyError:
+            logger.debug("No %s block", key)
+            continue
+
+        try:
+            pop = block.pop("$ext")
+        except (KeyError, AttributeError, TypeError):
+            logger.debug("No ext functions in %s block", key)
+            continue
+
+        func = get_wrapped_create_function(pop)
+        new_args[key] = func()
 
     merged_args = deep_dict_merge(request_args, new_args)
 

tavern/_plugins/mqtt/request.py

@@ -1,7 +1,7 @@
 import functools
 import json
 import logging
-from typing import Mapping
+from typing import Dict
 
 from box.box import Box
 
@@ -16,21 +16,19 @@
 logger = logging.getLogger(__name__)
 
 
-def get_publish_args(rspec: Mapping, test_block_config: TestConfig) -> dict:
-    """Format mqtt request args
-
-    Todo:
-        Anything else to do here?
-    """
+def get_publish_args(rspec: Dict, test_block_config: TestConfig) -> dict:
+    """Format mqtt request args and update using ext functions"""
 
     fspec = format_keys(rspec, test_block_config.variables)
 
-    if "json" in rspec:
-        if "payload" in rspec:
+    if "json" in fspec:
+        if "payload" in fspec:
             raise exceptions.BadSchemaError(
                 "Can only specify one of 'payload' or 'json' in MQTT request"
             )
 
+        update_from_ext(fspec, ["json"])
+
         fspec["payload"] = json.dumps(fspec.pop("json"))
 
     return fspec
@@ -43,15 +41,15 @@ class MQTTRequest(BaseRequest):
     """
 
     def __init__(
-        self, client: MQTTClient, rspec: Mapping, test_block_config: TestConfig
+        self, client: MQTTClient, rspec: Dict, test_block_config: TestConfig
     ) -> None:
         expected = {"topic", "payload", "json", "qos", "retain"}
 
         check_expected_keys(expected, rspec)
 
         publish_args = get_publish_args(rspec, test_block_config)
-        update_from_ext(publish_args, ["json"])
 
+        self._publish_args = publish_args
         self._prepared = functools.partial(client.publish, **publish_args)
 
         # Need to do this here because get_publish_args will modify the original

tavern/_plugins/mqtt/response.py

@@ -335,7 +335,7 @@ def _get_payload_vals(expected) -> Tuple[Optional[Union[str, dict]], bool]:
             json_payload = True
 
             if payload.pop("$ext", None):
-                raise exceptions.InvalidExtBlockException(
+                raise exceptions.MisplacedExtBlockException(
                     "json",
                 )
         elif "payload" in expected:

tavern/_plugins/rest/response.py

@@ -218,7 +218,7 @@ def _validate_block(self, blockname: str, block: Mapping) -> None:
 
         if isinstance(expected_block, dict):
             if expected_block.pop("$ext", None):
-                raise exceptions.InvalidExtBlockException(
+                raise exceptions.MisplacedExtBlockException(
                     blockname,
                 )
 

tavern/response.py

@@ -139,7 +139,7 @@ def check_deprecated_validate(name):
             if isinstance(block, dict):
                 check_ext_functions(block.get("$ext", None))
                 if nfuncs != len(self.validate_functions):
-                    raise exceptions.InvalidExtBlockException(
+                    raise exceptions.MisplacedExtBlockException(
                         name,
                     )
 

tests/unit/test_mqtt.py

@@ -1,3 +1,4 @@
+from typing import Dict
 from unittest.mock import MagicMock, Mock, patch
 
 import paho.mqtt.client as paho
@@ -18,18 +19,19 @@ def test_host_required():
     MQTTClient(**args)
 
 
-class TestClient:
-    @pytest.fixture(name="fake_client")
-    def fix_fake_client(self):
-        args = {"connect": {"host": "localhost"}}
+@pytest.fixture(name="fake_client")
+def fix_fake_client():
+    args = {"connect": {"host": "localhost"}}
+
+    mqtt_client = MQTTClient(**args)
 
-        mqtt_client = MQTTClient(**args)
+    mqtt_client._subscribed[2] = _Subscription("abc")
+    mqtt_client._subscription_mappings["abc"] = 2
 
-        mqtt_client._subscribed[2] = _Subscription("abc")
-        mqtt_client._subscription_mappings["abc"] = 2
+    return mqtt_client
 
-        return mqtt_client
 
+class TestClient:
     def test_no_queue(self, fake_client):
         """Trying to fetch from a nonexistent queue raised exception"""
 
@@ -192,3 +194,33 @@ def subscribe_success(topic, *args, **kwargs):
         MQTTClient._on_subscribe(mock_client, "abc", {}, 123, 0)
 
         assert mock_client._subscribed == {}
+
+
+class TestExtFunctions:
+    @pytest.fixture()
+    def basic_mqtt_request_args(self) -> Dict:
+        return {
+            "topic": "/a/b/c",
+        }
+
+    def test_basic(self, fake_client, basic_mqtt_request_args, includes):
+        MQTTRequest(fake_client, basic_mqtt_request_args, includes)
+
+    def test_ext_function_bad(self, fake_client, basic_mqtt_request_args, includes):
+        basic_mqtt_request_args["json"] = {"$ext": "kk"}
+
+        with pytest.raises(exceptions.InvalidExtFunctionError):
+            MQTTRequest(fake_client, basic_mqtt_request_args, includes)
+
+    def test_ext_function_good(self, fake_client, basic_mqtt_request_args, includes):
+        basic_mqtt_request_args["json"] = {
+            "$ext": {
+                "function": "operator:add",
+                "extra_args": (1, 2),
+            }
+        }
+
+        m = MQTTRequest(fake_client, basic_mqtt_request_args, includes)
+
+        assert "payload" in m._publish_args
+        assert m._publish_args["payload"] == "3"