Adjust test case command

This now correctly handles cases where there was no failure in the action.

When an action fails we save the initial state that went in. However, when
an action succeeds, we would only store the final state for that action.
So now, if the action was successful, we go back to the prior sequence
ID to get its state as the input, and use the state from the action as the
desired final state. This assumes you want to say "hey here's expected
behavior, I want to enshrine it".

Adjusts docs with new option.

burr/cli/__main__.py

@@ -244,19 +244,28 @@ def generate_demo_data(s3_bucket, data_dir, unique_app_names: bool, no_clear_cur
         generate_all(s3_bucket=s3_bucket, unique_app_names=unique_app_names)
 
 
-def _transform_state_to_test_case(state: dict, action_name: str, test_name: str) -> dict:
+def _remove_private_keys(state: dict):
+    return {key: value for key, value in state.items() if not key.startswith("__")}
+
+
+def _transform_state_to_test_case(
+    before_action_state: dict, after_action_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:
     """
+    # Remove private Burr keys -- don't want to expose these in the test case since we could change them
+    before_action_state = _remove_private_keys(before_action_state)
+    after_action_state = _remove_private_keys(after_action_state)
     return {
         "action": action_name,
         "name": test_name,
-        "input_state": state,
-        "expected_state": {"TODO:": "fill this in"},
+        "input_state": before_action_state,
+        "expected_state": after_action_state,
     }
 
 
@@ -308,8 +317,8 @@ def import_from_file(file_path: str) -> ModuleType:
 
 
 @click.command()
-@click.option("--project-name", required=True, help="Name of the project")
-@click.option("--partition-key", required=False, help="Partition key to look at")
+@click.option("--project-name", required=True, help="Name of the project.")
+@click.option("--partition-key", required=False, 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(
@@ -322,13 +331,19 @@ def import_from_file(file_path: str) -> ModuleType:
     required=False,
     help="Python file or fully qualified python module to import for custom serialization/deserialization.",
 )
+@click.option(
+    "--action-name",
+    required=False,
+    help="Provide name of the action for the test case. Defaults to the action name at the sequence ID.",
+)
 def create_test_case(
     project_name: str,
     partition_key: str,  # will be None if not passed in.
     app_id: str,
     sequence_id: str,
     target_file_name: str = None,
     serde_module: str = None,
+    action_name: str = None,
 ):
     """Create a test case from a persisted state.
 
@@ -343,6 +358,14 @@ def create_test_case(
     If you have custom serialization/deserialization then pass in the name of a
     python module, or a path to a python file, to import with your custom serialization/deserialization
     registration functions. This will be imported so that they can be registered.
+
+    Note:
+        - when burr fails, then the state retrieved for a sequence ID is the state at the start of the sequence ID.
+          The fixture created will only have input state.
+        - when burr completes successfully, then the state retrieved for a sequence ID is the final modified state.
+          The fixture created will have input and output state, as we will grab the prior sequence ID's state
+          as the input state.
+        - we strip any keys prefixed with __ from the state to avoid exposing private keys in the test case.
     """
     if serde_module:
         if os.path.exists(serde_module):
@@ -358,17 +381,60 @@ def create_test_case(
     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)
+    sequence_id = int(sequence_id)
+
+    after_action: PersistedStateData = local_tracker.load(
+        partition_key=partition_key, app_id=app_id, sequence_id=sequence_id
     )
-    if not data:
-        print(f"No data found for {app_id} in {project_name} with sequence {sequence_id}")
+    if not after_action:
+        print(f"No state found for {project_name}, {partition_key}, {app_id}, {sequence_id}.")
         return
-    state_dict = data["state"].serialize()
-    print("Found data for action: ", data["position"])
+
+    action_name = action_name if action_name else after_action["position"]
+    """
+    Explanation of the logic here.
+    If there's an error:
+     - status would be failed vs completed
+     - state loaded would be the starting state
+    If there's no error:
+     - status would be completed
+     - state loaded would be the ending state, so to get the starting state, we need to look at the prior sequence ID
+    """
+    if after_action["status"] == "completed":
+        print("Action was successful so loading initial and expected state into test fixture.")
+        after_action_state = after_action["state"].serialize()
+        # if it's completed, then we need to look at the prior sequence ID to get the input state
+        prior_sequence_id = sequence_id - 1
+        try:
+            # TODO: handle forked case. i.e. prior doesn't exist because it was a forked sequence.
+            before_action: PersistedStateData = local_tracker.load(
+                partition_key=partition_key, app_id=app_id, sequence_id=prior_sequence_id
+            )
+        except ValueError:
+            before_action = None
+
+        if not before_action:
+            # there was no initial state saved for this sequence id
+            print(
+                f"Warning: there was no initial state found sequence ID {sequence_id}. That is, we looked for the "
+                f"prior state corresponding to the prior sequence ID {prior_sequence_id}, but did not find a state "
+                f"value. Was this application ID forked? Defaulting to empty state - please fill this in."
+            )
+            before_action_state = {"TODO:": "fill this in"}
+        else:
+            before_action_state = before_action["state"].serialize()
+    else:
+        print("Action was not successful so loading initial state into test fixture.")
+        # there was an error so state is the starting state
+        before_action_state = after_action["state"].serialize()
+        after_action_state = {"TODO:": "fill this in"}
+
     # test case json
     tc_json = _transform_state_to_test_case(
-        state_dict, data["position"], f"{data['position']}_{app_id[:8] + '_' + str(sequence_id)}"
+        before_action_state,
+        after_action_state,
+        action_name,
+        f"{action_name}_{app_id[:8] + '_' + str(sequence_id)}",
     )
 
     if target_file_name:
@@ -387,7 +453,7 @@ def create_test_case(
         logger.info(json.dumps(tc_json, indent=2))
     # print out python test to add
     print("\nAdd the following to your test file:\n")
-    print(PYTEST_TEMPLATE.format(FILE_NAME=target_file_name, ACTION_NAME=data["position"]))
+    print(PYTEST_TEMPLATE.format(FILE_NAME=target_file_name, ACTION_NAME=action_name))
 
 
 test_case.add_command(create_test_case, name="create")

docs/examples/guardrails/creating_tests.rst

@@ -51,9 +51,12 @@ Steps:
 See `github repository example <https://github.com/DAGWorks-Inc/burr/tree/main/examples/test-case-creation>`_
 for an example.
 
-Note that, if you have custom serialization/deserialization logic, you will want to pass in `--serde-module` to the
+Note (1): if you have custom serialization/deserialization logic, you will want to pass in `--serde-module` to the
 test case with the module name of your serialization logic.
 
+Note (2): you can pass in `--action-name` to override the action name in the test case. This is useful if you want
+to use the output of one action as the input to another action; there are corner cases where this is useful.
+
 Future Work
 -----------
 We see many more improvements here:

examples/test-case-creation/notebook.ipynb

@@ -12,18 +12,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
    "id": "5cc77bd80999b2a2",
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-04-01T05:29:38.960440Z",
-     "start_time": "2024-04-01T05:29:37.926490Z"
-    },
     "collapsed": false,
     "jupyter": {
      "outputs_hidden": false
+    },
+    "ExecuteTime": {
+     "end_time": "2024-08-31T04:30:47.929817Z",
+     "start_time": "2024-08-31T04:30:46.259173Z"
     }
    },
+   "source": [
+    "%%sh \n",
+    "burr-test-case create --help"
+   ],
    "outputs": [
     {
      "name": "stdout",
@@ -41,21 +44,36 @@
       "  See examples/test-case-creation/notebook.ipynb for example usage. See\n",
       "  examples/test-case-creation/test_application.py for details.\n",
       "\n",
+      "  If you have custom serialization/deserialization then pass in the name of a\n",
+      "  python module, or a path to a python file, to import with your custom\n",
+      "  serialization/deserialization registration functions. This will be imported\n",
+      "  so that they can be registered.\n",
+      "\n",
+      "  Note:     - when burr fails, then the state retrieved for a sequence ID is\n",
+      "  the state at the start of the sequence ID.       The fixture created will\n",
+      "  only have input state.     - when burr completes successfully, then the\n",
+      "  state retrieved for a sequence ID is the final modified state.       The\n",
+      "  fixture created will have input and output state, as we will grab the prior\n",
+      "  sequence ID's state       as the input state.     - we strip any keys\n",
+      "  prefixed with __ from the state to avoid exposing private keys in the test\n",
+      "  case.\n",
+      "\n",
       "Options:\n",
-      "  --project-name TEXT      Name of the project  [required]\n",
-      "  --partition-key TEXT     Partition key to look at  [required]\n",
+      "  --project-name TEXT      Name of the project.  [required]\n",
+      "  --partition-key TEXT     Partition key to look at.\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",
+      "  --serde-module TEXT      Python file or fully qualified python module to\n",
+      "                           import for custom serialization/deserialization.\n",
+      "  --action-name TEXT       Provide name of the action for the test case.\n",
+      "                           Defaults to the action name at the sequence ID.\n",
       "  --help                   Show this message and exit.\n"
      ]
     }
    ],
-   "source": [
-    "%%sh \n",
-    "burr-test-case create --help"
-   ]
+   "execution_count": 1
   },
   {
    "cell_type": "markdown",