Repository: taverntesting/tavern

Issue:
Unable to use external function in MQTT publish
I tried using an external function in a MQTT publish request, and the function wasn't getting evaluated.

My request looks like this:
```
  - id: publish_thing_1
    name: Publish thing 1
    mqtt_publish:
      topic: &ping_topic '/device/123/ping'
      qos: 1
      json:
        $ext: 
          function: utils.testing_utils:my_function
        thing_1: abc
    mqtt_response:
      topic: *ping_topic
      json:
        header: test
        thing_1: abc
      timeout: 5
      qos: 1
```

It looks like the line below in [request.py](https://github.com/taverntesting/tavern/blob/master/tavern/_plugins/mqtt/request.py)
```
update_from_ext(publish_args, ["json"], test_block_config)
```
should be 
```
update_from_ext(publish_args, ["payload"], test_block_config)
```
instead, since publish_args looks like this:
```
{'topic': '/device/123/ping', 'qos': 1, 'payload': '{"$ext": {"function": "utils.testing_utils:my_function"}, "thing_1": "abc"}'}
```
Note that the `payload` value is a string, which prevents external function evaluation even after I do the `update_from_ext` change. Before I go down the rabbit hole too much, I wanted to confirm that I've configured the request properly, and that this feature is expected to work after this [PR](https://github.com/taverntesting/tavern/pull/620).

Comments:

Patch:
diff --git a/tavern/_core/exceptions.py b/tavern/_core/exceptions.py
index d8c78b6..803abe0 100644
--- a/tavern/_core/exceptions.py
+++ b/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:
diff --git a/tavern/_core/extfunctions.py b/tavern/_core/extfunctions.py
index c7b3325..e7b534e 100644
--- a/tavern/_core/extfunctions.py
+++ b/tavern/_core/extfunctions.py
@@ -123,6 +123,11 @@ def get_wrapped_create_function(ext: Mapping):
 
 
 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)
 
diff --git a/tavern/_plugins/mqtt/request.py b/tavern/_plugins/mqtt/request.py
index db7eae7..0a9de87 100644
--- a/tavern/_plugins/mqtt/request.py
+++ b/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 @@ from tavern.request import BaseRequest
 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
diff --git a/tavern/_plugins/mqtt/response.py b/tavern/_plugins/mqtt/response.py
index a2e362b..73513f8 100644
--- a/tavern/_plugins/mqtt/response.py
+++ b/tavern/_plugins/mqtt/response.py
@@ -335,7 +335,7 @@ class _MessageVerifier:
             json_payload = True
 
             if payload.pop("$ext", None):
-                raise exceptions.InvalidExtBlockException(
+                raise exceptions.MisplacedExtBlockException(
                     "json",
                 )
         elif "payload" in expected:
diff --git a/tavern/_plugins/rest/response.py b/tavern/_plugins/rest/response.py
index 97bc494..ca54b11 100644
--- a/tavern/_plugins/rest/response.py
+++ b/tavern/_plugins/rest/response.py
@@ -218,7 +218,7 @@ class RestResponse(BaseResponse):
 
         if isinstance(expected_block, dict):
             if expected_block.pop("$ext", None):
-                raise exceptions.InvalidExtBlockException(
+                raise exceptions.MisplacedExtBlockException(
                     blockname,
                 )
 
diff --git a/tavern/response.py b/tavern/response.py
index f9eba9f..8bee7a6 100644
--- a/tavern/response.py
+++ b/tavern/response.py
@@ -139,7 +139,7 @@ class BaseResponse:
             if isinstance(block, dict):
                 check_ext_functions(block.get("$ext", None))
                 if nfuncs != len(self.validate_functions):
-                    raise exceptions.InvalidExtBlockException(
+                    raise exceptions.MisplacedExtBlockException(
                         name,
                     )
 


Test patch:
diff --git a/example/mqtt/test_mqtt.tavern.yaml b/example/mqtt/test_mqtt.tavern.yaml
index 145d3b6..956a18a 100644
--- a/example/mqtt/test_mqtt.tavern.yaml
+++ b/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
diff --git a/example/mqtt/testing_utils.py b/example/mqtt/testing_utils.py
index 70021ab..f483ca9 100644
--- a/example/mqtt/testing_utils.py
+++ b/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"}
diff --git a/tests/unit/test_mqtt.py b/tests/unit/test_mqtt.py
index b9603da..1d55b06 100644
--- a/tests/unit/test_mqtt.py
+++ b/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 @@ class TestSubscription:
         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"