Files
extism/python/tests/test_extism.py
zach 94d5bd98c8 feat: Add plugin cancellation (#270)
This PR adds the ability to cancel running plugins from another thread
in the host.

- All SDKs have been updated except  .NET
- Adds a new SDK type: `ExtismCancelHandle`
- Adds 2 new SDK functions: `extism_plugin_cancel_handle` and
`extism_plugin_cancel`
- `extism_plugin_cancel_handle` returns a pointer to
`ExtismCancelHandle`, which can be passed to `extism_plugin_cancel` to
stop plugin execution.
- One thing that's worth noting is when plugin is executing a host
function it cannot be cancelled until after the host function returns.

---------

Co-authored-by: Etienne ANNE <etienne.anne@icloud.com>
2023-03-08 17:44:40 -08:00

138 lines
4.9 KiB
Python

import unittest
import extism
import hashlib
import json
import time
from threading import Thread
from datetime import datetime, timedelta
from os.path import join, dirname
class TestExtism(unittest.TestCase):
def test_context_new(self):
ctx = extism.Context()
self.assertIsNotNone(ctx)
del ctx
def test_call_plugin(self):
with extism.Context() as ctx:
plugin = ctx.plugin(self._manifest())
j = json.loads(plugin.call("count_vowels", "this is a test"))
self.assertEqual(j["count"], 4)
j = json.loads(plugin.call("count_vowels", "this is a test again"))
self.assertEqual(j["count"], 7)
j = json.loads(plugin.call("count_vowels", "this is a test thrice"))
self.assertEqual(j["count"], 6)
j = json.loads(plugin.call("count_vowels", "🌎hello🌎world🌎"))
self.assertEqual(j["count"], 3)
def test_update_plugin_manifest(self):
with extism.Context() as ctx:
plugin = ctx.plugin(self._manifest())
# update with just the raw bytes of the wasm
plugin.update(self._count_vowels_wasm())
# should still work
j = json.loads(plugin.call("count_vowels", "this is a test"))
self.assertEqual(j["count"], 4)
def test_function_exists(self):
with extism.Context() as ctx:
plugin = ctx.plugin(self._manifest())
self.assertTrue(plugin.function_exists("count_vowels"))
self.assertFalse(plugin.function_exists("i_dont_exist"))
def test_errors_on_unknown_function(self):
with extism.Context() as ctx:
plugin = ctx.plugin(self._manifest())
self.assertRaises(
extism.Error, lambda: plugin.call("i_dont_exist", "someinput")
)
def test_can_free_plugin(self):
with extism.Context() as ctx:
plugin = ctx.plugin(self._manifest())
del plugin
def test_errors_on_bad_manifest(self):
with extism.Context() as ctx:
self.assertRaises(
extism.Error, lambda: ctx.plugin({"invalid_manifest": True})
)
plugin = ctx.plugin(self._manifest())
self.assertRaises(
extism.Error, lambda: plugin.update({"invalid_manifest": True})
)
def test_extism_version(self):
self.assertIsNotNone(extism.extism_version())
def test_extism_plugin_timeout(self):
with extism.Context() as ctx:
plugin = ctx.plugin(self._loop_manifest())
start = datetime.now()
self.assertRaises(extism.Error, lambda: plugin.call("infinite_loop", b""))
end = datetime.now()
self.assertLess(
end,
start + timedelta(seconds=1.01),
"plugin timeout exceeded 1000ms expectation",
)
def test_extism_host_function(self):
@extism.host_fn
def hello_world(plugin, params, results, user_data):
offs = plugin.alloc(len(user_data))
mem = plugin.memory(offs)
mem[:] = user_data
results[0].value = offs.offset
with extism.Context() as ctx:
f = [
extism.Function(
"hello_world",
[extism.ValType.I64],
[extism.ValType.I64],
hello_world,
b"test",
)
]
plugin = ctx.plugin(self._manifest(functions=True), functions=f, wasi=True)
res = plugin.call("count_vowels", "aaa")
self.assertEqual(res, b"test")
def test_extism_plugin_cancel(self):
with extism.Context() as ctx:
plugin = ctx.plugin(self._loop_manifest())
cancel_handle = plugin.cancel_handle()
def cancel(handle):
time.sleep(0.5)
handle.cancel()
Thread(target=cancel, args=[cancel_handle]).run()
self.assertRaises(extism.Error, lambda: plugin.call("infinite_loop", b""))
def _manifest(self, functions=False):
wasm = self._count_vowels_wasm(functions)
hash = hashlib.sha256(wasm).hexdigest()
return {"wasm": [{"data": wasm, "hash": hash}], "memory": {"max_pages": 5}}
def _loop_manifest(self):
wasm = self._infinite_loop_wasm()
hash = hashlib.sha256(wasm).hexdigest()
return {
"wasm": [{"data": wasm, "hash": hash}],
"memory": {"max_pages": 5},
"timeout_ms": 1000,
}
def _count_vowels_wasm(self, functions=False):
return read_test_wasm("code.wasm" if not functions else "code-functions.wasm")
def _infinite_loop_wasm(self):
return read_test_wasm("loop.wasm")
def read_test_wasm(p):
path = join(dirname(__file__), "..", "..", "wasm", p)
with open(path, "rb") as wasm_file:
return wasm_file.read()