Add pytest (#16)

.github/workflows/autoblocks-simulations.yml

@@ -9,6 +9,12 @@ env:
   POETRY_VERSION: "1.5.1"
   PYTHON_VERSION: "3.11"
 
+  # Use the simulation ingestion key so that Autoblocks knows we're sending simulated events
+  AUTOBLOCKS_INGESTION_KEY: ${{ secrets.AUTOBLOCKS_SIMULATION_INGESTION_KEY }}
+
+  # Any other environment variables the application needs to run
+  OPENAI_API_KEY: ${{ secrets.DEMO_OPENAI_API_KEY }}
+
 jobs:
   autoblocks-simulations:
     runs-on: ubuntu-latest
@@ -32,25 +38,19 @@ jobs:
       - name: Install dependencies
         run: poetry install
 
+      - name: Run tests
+        run: poetry run pytest
+
       - name: Start the app
         run: poetry run start &
-        env:
-          # Use the simulation ingestion key so that Autoblocks knows we're sending simulated events
-          AUTOBLOCKS_INGESTION_KEY: ${{ secrets.AUTOBLOCKS_SIMULATION_INGESTION_KEY }}
-
-          # Any other environment variables the application needs to run
-          OPENAI_API_KEY: ${{ secrets.DEMO_OPENAI_API_KEY }}
 
       - name: Wait for the app to be ready
         run: |
           while [[ "$(curl -s -o /dev/null -w ''%{http_code}'' http://localhost:5000/health)" != "200" ]]; do sleep 1; done
 
-      - name: Run simulation with static test cases
-        run: poetry run simulation-static
-
       - name: Run simulation with production test cases
         run: poetry run simulation-production-replay
         env:
           # Production events are fetched from the Autoblocks API,
-          # so we need the API key to authenticate
+          # so we need the API key to authenticate.
           AUTOBLOCKS_API_KEY: ${{ secrets.AUTOBLOCKS_API_KEY }}

README.md

@@ -4,39 +4,42 @@ This repository contains an example project that integrates Autoblocks Simulatio
 via GitHub Actions.
 See the [documentation](https://docs.autoblocks.ai/guides/simulations) for more information.
 
-## Instructions
+## Instructions for running locally
 
 ### 1. Install dependencies
 
 ```bash
 poetry install
 ```
 
-### 2. Start the application
+### 2. Run the tests
 
-Start the application with your simulation ingestion key,
-a simulation id to uniquely identify your simulation run,
-and any other environment variables needed to run your application:
+Set the `AUTOBLOCKS_INGESTION_KEY` environment variable to your simulation ingestion key before running the tests
+so that any events sent during the test run are sent as simulated events.
 
 ```bash
 AUTOBLOCKS_INGESTION_KEY=<simulation-ingestion-key> \
-AUTOBLOCKS_SIMULATION_ID=$(date +%Y%m%d%H%M%S) \
 OPENAI_API_KEY=<openai-api-key> \
-poetry run start
+poetry run pytest
 ```
 
-### 3. Run the simulation
+[View your simulation](https://app.autoblocks.ai/simulations)
 
-In another terminal, run either:
+### 3. Run a simulation with production events
 
-* `simulation-static`, which will replay a static set of test cases against your application:
+To run a simulation that replays production events, first start the application:
 
 ```bash
-poetry run simulation-static
+AUTOBLOCKS_INGESTION_KEY=<simulation-ingestion-key> \
+AUTOBLOCKS_SIMULATION_ID=$(date +%Y%m%d%H%M%S) \
+OPENAI_API_KEY=<openai-api-key> \
+poetry run start
 ```
 
-* `simulation-production-replay`, which will replay a set of production events fetched from the Autoblocks API:
+Then, in a separate terminal, run the simulation:
 
 ```bash
 AUTOBLOCKS_API_KEY=<autoblocks-api-key> poetry run simulation-production-replay
 ```
+
+[View your simulation](https://app.autoblocks.ai/simulations)

demo_app/app.py

@@ -1,13 +1,10 @@
 import uuid
 
-from autoblocks.tracer import AutoblocksTracer
 from flask import Flask
 from flask import request
 
 from demo_app import bot
 from demo_app.settings import AUTOBLOCKS_SIMULATION_TRACE_ID_HEADER_NAME
-from demo_app.settings import REQUEST_PAYLOAD_MESSAGE
-from demo_app.settings import env
 
 app = Flask(__name__)
 
@@ -29,19 +26,9 @@ def main():
     # but in a simulation scenario we use the trace id passed in via the simulation trace id header
     trace_id = request.headers.get(AUTOBLOCKS_SIMULATION_TRACE_ID_HEADER_NAME) or str(uuid.uuid4())
 
-    autoblocks = AutoblocksTracer(
-        env.AUTOBLOCKS_INGESTION_KEY,
-        trace_id=trace_id,
-        properties=dict(source="DEMO_SIMULATIONS"),
-    )
-    autoblocks.send_event(REQUEST_PAYLOAD_MESSAGE, properties=dict(payload=payload))
+    output = bot.get_response(trace_id, query)
 
-    output = bot.get_response(autoblocks, query)
-
-    response = {"output": output}
-    autoblocks.send_event("request.response", properties=dict(response=response))
-
-    return response
+    return dict(output=output)
 
 
 def start():

demo_app/bot.py

@@ -6,7 +6,13 @@
 from demo_app.settings import env
 
 
-def get_response(autoblocks: AutoblocksTracer, query: str) -> str:
+def get_response(trace_id: str, query: str) -> str:
+    autoblocks = AutoblocksTracer(
+        env.AUTOBLOCKS_INGESTION_KEY,
+        trace_id=trace_id,
+        properties=dict(source="DEMO_SIMULATIONS"),
+    )
+
     ai = AIChat(
         api_key=env.OPENAI_API_KEY,
         model="gpt-3.5-turbo",

demo_app/settings.py

@@ -4,9 +4,9 @@
 # the event being handled during the simulation.
 AUTOBLOCKS_SIMULATION_TRACE_ID_HEADER_NAME = "x-autoblocks-simulation-trace-id"
 
-# The message for the request.payload event, pulled into a variable here
+# The message for the user input event, pulled into a variable here
 # so that it's kept in sync between sending the events and replaying them.
-REQUEST_PAYLOAD_MESSAGE = "request.payload"
+USER_QUERY_MESSAGE = "user.query"
 
 
 # Environment variables

demo_app/simulations.py

@@ -8,28 +8,10 @@
 from autoblocks.api.models import TraceFilterOperator
 
 from demo_app.settings import AUTOBLOCKS_SIMULATION_TRACE_ID_HEADER_NAME
-from demo_app.settings import REQUEST_PAYLOAD_MESSAGE
+from demo_app.settings import USER_QUERY_MESSAGE
 from demo_app.settings import env
 
 
-def static():
-    """
-    Test a static set of events against the locally-running app.
-    """
-    for trace_id, query in [
-        ("san-francisco-tourist-attractions", "San Francisco tourist attractions"),
-        ("paris-tourist-attractions", "Paris tourist attractions"),
-        ("lombard-street", "Lombard Street"),
-        ("eiffel-tower", "Eiffel Tower"),
-    ]:
-        print(f"Testing static event {trace_id} - {query}")
-        requests.post(
-            "http://localhost:5000",
-            json={"query": query},
-            headers={AUTOBLOCKS_SIMULATION_TRACE_ID_HEADER_NAME: trace_id},
-        )
-
-
 def production_replay():
     """
     Replays production events fetched from the Autoblocks API against the locally-running app.
@@ -46,15 +28,15 @@ def production_replay():
                     EventFilter(
                         key=SystemEventFilterKey.MESSAGE,
                         operator=EventFilterOperator.EQUALS,
-                        value=REQUEST_PAYLOAD_MESSAGE,
+                        value=USER_QUERY_MESSAGE,
                     ),
                 ],
             ),
         ],
     )
     for trace in page.traces:
         for event in trace.events:
-            if event.message == REQUEST_PAYLOAD_MESSAGE:
+            if event.message == USER_QUERY_MESSAGE:
                 print(f"Replaying past event {event}")
 
                 # The original payload

pyproject.toml

@@ -13,10 +13,11 @@ requests = "^2.31.0"
 flask = "^2.3.2"
 pydantic-settings = "^2.0.2"
 simpleaichat = "^0.2.2"
-autoblocksai = "0.0.6"
+autoblocksai = "0.0.7"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.3.3"
+pytest = "^7.4.0"
 
 [build-system]
 requires = ["poetry-core"]
@@ -35,5 +36,4 @@ known-first-party = ["demo_app"]
 
 [tool.poetry.scripts]
 start = "demo_app.app:start"
-simulation-static = "demo_app.simulations:static"
 simulation-production-replay = "demo_app.simulations:production_replay"

tests/conftest.py

@@ -0,0 +1,11 @@
+import os
+from datetime import datetime
+
+import pytest
+
+
+@pytest.fixture(scope="session", autouse=True)
+def set_autoblocks_simulation_id():
+    os.environ["AUTOBLOCKS_SIMULATION_ID"] = "pytest-" + datetime.now().strftime("%Y%m%d-%H%M%S")
+    yield
+    del os.environ["AUTOBLOCKS_SIMULATION_ID"]

tests/test_bot.py

@@ -0,0 +1,17 @@
+import pytest
+
+from demo_app import bot
+
+
+@pytest.mark.parametrize(
+    "trace_id,query,expected_output",
+    [
+        ("san-francisco-tourist-attractions", "San Francisco tourist attractions", "Lombard"),
+        ("paris-tourist-attractions", "Paris tourist attractions", "Eiffel"),
+        ("lombard-street", "Lombard Street", "San Francisco"),
+        ("eiffel-tower", "Eiffel Tower", "Paris"),
+    ],
+)
+def test_bot(trace_id: str, query: str, expected_output: str):
+    response = bot.get_response(trace_id, query)
+    assert expected_output in response