Adds command to create a test case

The idea is that someone is in the tracker and wants to fix/create
a test around some part of the application.

They can then use the command line to bootstrap it. We could
do this in the UI too eventually.

We assume they want to then have a JSON file with input state
and expected output state to validate against. I think this
is a nice way to do it. Then it's all checked in and they can run
it in CI easily...

This command assumes a local tracker for now. In the future
we can add a yaml/config file that can instantiate a persister
to load data instead.

Otherwise we assume pytest, we could create a unittest stub too,
and also that users will like the file drive approach.

burr/cli/__main__.py

@@ -1,3 +1,4 @@
+import json
 import os
 import shutil
 import subprocess
@@ -9,6 +10,7 @@
 from importlib.resources import files
 
 from burr import telemetry
+from burr.core.persistence import PersistedStateData
 from burr.integrations.base import require_plugin
 
 try:
@@ -175,10 +177,109 @@ def generate_demo_data():
         generate_all("burr/tracking/server/demo_data")
 
 
+def _transform_state_to_test_case(state: dict, action_name: str, test_name: str) -> dict:
+    """Helper function to transform a state into a test case.
+
+    :param state:
+    :param action_name:
+    :param test_name:
+    :return:
+    """
+    return {
+        "action": action_name,
+        "name": test_name,
+        "input_state": state,
+        "expected_state": {"TODO:": "fill this in"},
+    }
+
+
+@click.group()
+def test_case():
+    """Test case related commands."""
+    pass
+
+
+PYTEST_TEMPLATE = """import pytest
+# TODO: import the action you're testing
+@pytest.mark.file_name("{FILE_NAME}")
+def test_{ACTION_NAME}(input_state, expected_state):
+    \"\"\"Function for testing the action\"\"\"
+    input_state = state.State(input_state)
+    expected_state = state.State(expected_state)
+    _, output_state = {ACTION_NAME}(input_state)  # exercise the action
+    assert output_state == expected_state
+"""
+
+
+@click.command()
+@click.option("--project-name", required=True, help="Name of the project")
+@click.option("--partition-key", required=True, help="Partition key to look at")
+@click.option("--app-id", required=True, help="App ID to pull from.")
+@click.option("--sequence-id", required=True, help="Sequence ID to pull.")
+@click.option(
+    "--target-file-name",
+    required=False,
+    help="What file to write the data to. Else print to console.",
+)
+def create_test_case(
+    project_name: str,
+    partition_key: str,
+    app_id: str,
+    sequence_id: str,
+    target_file_name: str = None,
+):
+    """Create a test case from a persisted state.
+
+    Prints a pytest test case to the console assuming you have a `pytest_generate_tests` function set up.
+
+    See examples/test-case-creation/test_application.py for details.
+    """
+    # TODO: make this handle instantiating/using a persister other than local tracker
+    from burr.tracking.client import LocalTrackingClient
+
+    local_tracker = LocalTrackingClient(project=project_name)
+    data: PersistedStateData = local_tracker.load(
+        partition_key=partition_key, app_id=app_id, sequence_id=int(sequence_id)
+    )
+    if not data:
+        print(f"No data found for {app_id} in {project_name} with sequence {sequence_id}")
+        return
+    state_dict = data["state"].get_all()
+    print("Found data for action: ", data["position"])
+    # test case json
+    tc_json = _transform_state_to_test_case(
+        state_dict, data["position"], f"{data['position']}_{app_id[:8] + '_' + str(sequence_id)}"
+    )
+
+    if target_file_name:
+        # if it already exists, load it up and append to it
+        if os.path.exists(target_file_name):
+            with open(target_file_name, "r") as f:
+                # assumes it's a list of test cases
+                current_testcases = json.load(f)
+            current_testcases.append(tc_json)
+        else:
+            current_testcases = [tc_json]
+        print(f"\nWriting data to file {target_file_name}")
+        with open(target_file_name, "w") as f:
+            json.dump(current_testcases, f, indent=2)
+    else:
+        logger.info(json.dumps(tc_json, indent=2))
+    # print out python test to add
+    print(
+        "\nAdd the following to your test file assuming you've got `pytest_generate_tests` set up:\n"
+    )
+    print(PYTEST_TEMPLATE.format(FILE_NAME=target_file_name, ACTION_NAME=data["position"]))
+
+
+test_case.add_command(create_test_case, name="create")
+cli.add_command(test_case)
+
 # quick trick to expose every subcommand as a variable
 # will create a command called `cli_{command}` for every command we have
 for key, command in cli.commands.items():
     globals()[f'cli_{key.replace("-", "_")}'] = command
 
+
 if __name__ == "__main__":
     cli()

docs/examples/creating_tests.rst

@@ -0,0 +1,14 @@
+====================
+Creating Test Cases
+====================
+
+With Burr tracking state is part of the framework. This means creating a realistic test case
+for an action involves turning a persister or tracker on and then pulling that state out
+for a test case. The following example demonstrates how to create a test case
+using the `burr-test-case` command.
+
+
+Test Case Creation Example
+--------------------------
+
+See `github repository example <https://github.com/DAGWorks-Inc/burr/tree/main/examples/test-case-creation>`_.

docs/examples/index.rst

@@ -15,3 +15,4 @@ Examples of more complex/powerful use-cases of Burr. Download/copy these to adap
     agents
     ml_training
     simulation
+    creating_tests

examples/test_case_creation/README.md → examples/test-case-creation/README.md

@@ -17,7 +17,7 @@ real-world behavior of your application -- Burr can make this process easier.
 2. Run your Burr application with tracking/persistence.
 3. Find a trace that you want to create a test case for.
 4. Note the partition_key, app_id, and sequence_id of the trace.
-5. Create a test case using the `burr create-test-case` command passing the partition_key, app_id, and sequence_id as arguments.
+5. Create a test case using the `burr-test-case create` command passing the project name, partition_key, app_id, and sequence_id as arguments.
 6. Iterate on the test case until it is robust and representative of the real-world behavior of your application.
 7. Profit.
 

examples/test_case_creation/application.py → examples/test-case-creation/application.py

@@ -1,11 +1,12 @@
+"""This file is truncated to just the relevant parts for the example."""
 from typing import Tuple
 
 import openai
 
 from burr.core import State
 from burr.core.action import action
 
-#  This file is truncated to just the relevant parts for the example.
+# Reminder: This file is truncated to just the relevant parts for the example.
 
 MODES = {
     "answer_question": "text",
@@ -57,4 +58,4 @@ def prompt_for_more(state: State) -> Tuple[dict, State]:
     return result, state.update(**result)
 
 
-# This file is truncated to just the relevant parts for the example.
+# Reminder: this file is truncated to just the relevant parts for the example.

examples/test-case-creation/notebook.ipynb

@@ -0,0 +1,142 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Usage: burr-test-case create [OPTIONS]\n",
+      "\n",
+      "  Create a test case from a persisted state.\n",
+      "\n",
+      "  Prints a pytest test case to the console assuming you have a\n",
+      "  `pytest_generate_tests` function set up.\n",
+      "\n",
+      "  See examples/test-case-creation/test_application.py for details.\n",
+      "\n",
+      "Options:\n",
+      "  --project-name TEXT      Name of the project  [required]\n",
+      "  --partition-key TEXT     Partition key to look at  [required]\n",
+      "  --app-id TEXT            App ID to pull from.  [required]\n",
+      "  --sequence-id TEXT       Sequence ID to pull.  [required]\n",
+      "  --target-file-name TEXT  What file to write the data to. Else print to\n",
+      "                           console.\n",
+      "  --help                   Show this message and exit.\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%sh \n",
+    "burr-test-case create --help\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-04-01T05:29:38.960440Z",
+     "start_time": "2024-04-01T05:29:37.926490Z"
+    }
+   },
+   "id": "5cc77bd80999b2a2",
+   "execution_count": 7
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Create a test case\n",
+    "\n",
+    "Find the data you want to create a test case with in the UI."
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "944a988132d19779"
+  },
+  {
+   "cell_type": "code",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Found data for action:  plan_cot\n",
+      "\n",
+      "Writing data to file /tmp/test-case.json\n",
+      "\n",
+      "Add the following to your test file assuming you've got `pytest_generate_tests` set up:\n",
+      "\n",
+      "import pytest\n",
+      "# TODO: import the action you're testing\n",
+      "@pytest.mark.file_name(\"/tmp/test-case.json\")\n",
+      "def test_plan_cot(input_state, expected_state):\n",
+      "    \"\"\"Function for testing the action\"\"\"\n",
+      "    input_state = state.State(input_state)\n",
+      "    expected_state = state.State(expected_state)\n",
+      "    _, output_state = plan_cot(input_state)  # exercise the action\n",
+      "    assert output_state == expected_state\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%sh\n",
+    "burr-test-case create --project-name \"oa:determine_params\" --partition-key \"\" --app-id \"3350dd68-bcfd-48a1-85c5-dd72186382ef\" --sequence-id 0 --target-file-name /tmp/test-case.json\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-04-01T05:33:01.015032Z",
+     "start_time": "2024-04-01T05:33:00.146352Z"
+    }
+   },
+   "id": "586bcb0bdb5dde8b",
+   "execution_count": 11
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# What's next?\n",
+    " \n",
+    "## See `test_application.py` for code that enables you to run the test case like this.\n",
+    "\n",
+    "## Fill in the expected state part of the JSON\n",
+    "The expected state part of the JSON is empty. You need to fill it in with the expected state of the application after the test case has been run. Once you've\n",
+    "done that, you've got your test case!\n"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "3ac9d600d0807fc8"
+  },
+  {
+   "cell_type": "code",
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "a727cdc01303cd39"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -133,3 +133,4 @@ burr-admin-server = "burr.cli.__main__:cli_run_server"
 burr-admin-publish = "burr.cli.__main__:cli_build_and_publish"
 burr-admin-build-ui = "burr.cli.__main__:cli_build_ui"
 burr-admin-generate-demo-data = "burr.cli.__main__:cli_generate_demo_data"
+burr-test-case = "burr.cli.__main__:cli_test_case"