Merge pull request #126 from stealthrocket/test-server

Mock server

Author: chriso

README.md

@@ -22,8 +22,9 @@ Python package to develop applications with the Dispatch platform.
 - [Usage](#usage)
   - [Configuration](#configuration)
   - [Integration with FastAPI](#integration-with-fastapi)
-  - [Local testing with ngrok](#local-testing-with-ngrok)
-  - [Distributed coroutines for Python](#distributed-coroutines-for-python)
+  - [Local Testing](#local-testing)
+  - [Distributed Coroutines for Python](#distributed-coroutines-for-python)
+  - [Serialization](#serialization)
 - [Examples](#examples)
 - [Contributing](#contributing)
 
@@ -123,10 +124,46 @@ program, driven by the Dispatch SDK.
 The instantiation of the `Dispatch` object on the `FastAPI` application
 automatically installs the HTTP route needed for Dispatch to invoke functions.
 
-### Local testing with ngrok
+### Local Testing
 
-To enable local testing, a common approach consists of using [ngrok][ngrok] to
-setup a public endpoint that forwards to the server running on localhost.
+#### Mock Dispatch
+
+The SDK ships with a mock Dispatch server. It can be used to quickly test your
+local functions, without requiring internet access.
+
+Note that the mock Dispatch server has very limited scheduling capabilities.
+
+```console
+python -m dispatch.test $DISPATCH_ENDPOINT_URL
+```
+
+The command will start a mock Dispatch server and print the configuration
+for the SDK.
+
+For example, if your functions were exposed through a local endpoint
+listening on `http://127.0.0.1:8000`, you could run:
+
+```console
+$ python -m dispatch.test http://127.0.0.1:8000
+Spawned a mock Dispatch server on 127.0.0.1:4450
+
+Dispatching function calls to the endpoint at http://127.0.0.1:8000
+
+The Dispatch SDK can be configured with:
+
+  export DISPATCH_API_URL="http://127.0.0.1:4450"
+  export DISPATCH_API_KEY="test"
+  export DISPATCH_ENDPOINT_URL="http://127.0.0.1:8000"
+  export DISPATCH_VERIFICATION_KEY="Z+nTe2VRcw8t8Ihx++D+nXtbO28nwjWIOTLRgzrelYs="
+```
+
+#### Real Dispatch
+
+To test local functions with the production instance of Dispatch, it needs
+to be able to access your local endpoint.
+
+A common approach consists of using [ngrok][ngrok] to setup a public endpoint
+that forwards to the server running on localhost.
 
 For example, assuming the server is running on port 8000 (which is the default
 with FastAPI), the command to create a ngrok tunnel is:

pyproject.toml

@@ -30,6 +30,9 @@ dev = [
     "coverage >= 7.4.1",
     "requests >= 2.31.0",
     "types-requests >= 2.31.0.20240125",
+    "docopt >= 0.6.2",
+    "types-docopt >= 0.6.11.4",
+    "uvicorn >= 0.28.0"
 ]
 
 docs = [

src/dispatch/scheduler.py

@@ -329,8 +329,8 @@ def _run(self, input: Input) -> Output:
                     coroutine_id=coroutine.id, value=e.value
                 )
             except Exception as e:
-                logger.exception(
-                    f"@dispatch.function: '{coroutine}' raised an exception"
+                logger.debug(
+                    f"@dispatch.function: '{coroutine}' raised an exception", exc_info=e
                 )
                 coroutine_result = CoroutineResult(coroutine_id=coroutine.id, error=e)
 

src/dispatch/status.py

@@ -30,6 +30,12 @@ class Status(int, enum.Enum):
 
     _proto: status_pb.Status
 
+    def __repr__(self):
+        return self.name
+
+    def __str__(self):
+        return self.name
+
 
 # Maybe we should find a better way to define that enum. It's that way to please
 # Mypy and provide documentation for the enum values.

src/dispatch/test/__main__.py

@@ -0,0 +1,92 @@
+"""Mock Dispatch server for use in test environments.
+
+Usage:
+  dispatch.test <endpoint> [--api-key=<key>] [--hostname=<name>] [--port=<port>] [-v | --verbose]
+  dispatch.test -h | --help
+
+Options:
+     --api-key=<key>      API key to require when clients connect to the server [default: test].
+
+     --hostname=<name>    Hostname to listen on [default: 127.0.0.1].
+     --port=<port>        Port to listen on [default: 4450].
+
+  -v --verbose            Show verbose details in the log.
+  -h --help               Show this help information.
+"""
+
+import base64
+import logging
+import os
+import sys
+
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
+from docopt import docopt
+
+from dispatch.test import DispatchServer, DispatchService, EndpointClient
+
+
+def main():
+    args = docopt(__doc__)
+
+    if args["--help"]:
+        print(__doc__)
+        exit(0)
+
+    endpoint = args["<endpoint>"]
+    api_key = args["--api-key"]
+    hostname = args["--hostname"]
+    port_str = args["--port"]
+
+    try:
+        port = int(port_str)
+    except ValueError:
+        print(f"error: invalid port: {port_str}", file=sys.stderr)
+        exit(1)
+
+    if not os.getenv("NO_COLOR"):
+        logging.addLevelName(logging.WARNING, f"\033[1;33mWARN\033[1;0m")
+        logging.addLevelName(logging.ERROR, "\033[1;31mERROR\033[1;0m")
+
+    logger = logging.getLogger()
+    if args["--verbose"]:
+        logger.setLevel(logging.DEBUG)
+        fmt = "%(asctime)s [%(levelname)s] %(name)s - %(message)s"
+    else:
+        logger.setLevel(logging.INFO)
+        fmt = "%(asctime)s [%(levelname)s] %(message)s"
+        logging.getLogger("httpx").disabled = True
+
+    log_formatter = logging.Formatter(fmt=fmt, datefmt="%Y-%m-%d %H:%M:%S")
+    log_handler = logging.StreamHandler(sys.stderr)
+    log_handler.setFormatter(log_formatter)
+    logger.addHandler(log_handler)
+
+    # This private key was generated randomly.
+    signing_key = Ed25519PrivateKey.from_private_bytes(
+        b"\x0e\xca\xfb\xc9\xa9Gc'fR\xe4\x97y\xf0\xae\x90\x01\xe8\xd9\x94\xa6\xd4@\xf6\xa7!\x90b\\!z!"
+    )
+    verification_key = base64.b64encode(
+        signing_key.public_key().public_bytes_raw()
+    ).decode()
+
+    endpoint_client = EndpointClient.from_url(endpoint, signing_key=signing_key)
+
+    with DispatchService(endpoint_client, api_key=api_key) as service:
+        with DispatchServer(service, hostname=hostname, port=port) as server:
+            print(f"Spawned a mock Dispatch server on {hostname}:{port}")
+            print()
+            print(f"Dispatching function calls to the endpoint at {endpoint}")
+            print()
+            print("The Dispatch SDK can be configured with:")
+            print()
+            print(f'  export DISPATCH_API_URL="http://{hostname}:{port}"')
+            print(f'  export DISPATCH_API_KEY="{api_key}"')
+            print(f'  export DISPATCH_ENDPOINT_URL="{endpoint}"')
+            print(f'  export DISPATCH_VERIFICATION_KEY="{verification_key}"')
+            print()
+
+            server.wait()
+
+
+if __name__ == "__main__":
+    main()

src/dispatch/test/server.py

@@ -39,6 +39,10 @@ def start(self):
         """Start the server."""
         self._server.start()
 
+    def wait(self):
+        """Block until the server terminates."""
+        self._server.wait_for_termination()
+
     def stop(self):
         """Stop the server."""
         self._server.stop(0)