Merge pull request #819 from PrefectHQ/call-routing-fix-tests

fix tests

docs/examples/call_routing.md

@@ -31,7 +31,7 @@ Automatically route customer calls to the right department.
 
     !!! success "💵 Price matching"
         ```python
-        department = route_call("Do you price match?")
+        department = route_call("Well FooCo offered me a better deal")
         assert department == Department.SALES
         ```
 

tests/ai/test_classify.py

@@ -2,6 +2,7 @@
 from typing import Literal
 
 import marvin
+import pytest
 
 Sentiment = Literal["Positive", "Negative"]
 
@@ -83,11 +84,15 @@ async def test_hogwarts_sorting_hat(self):
 
             assert house == "Gryffindor"
 
-        async def test_call_routing(self):
-            from enum import Enum
-
-            import marvin
-
+        @pytest.mark.parametrize(
+            "user_input, expected_selection",
+            [
+                ("I need to update my payment method", "billing"),
+                ("Well FooCo offered me a better deal", "sales"),
+                ("*angry noises*", "support"),
+            ],
+        )
+        async def test_call_routing(self, user_input, expected_selection):
             class Department(Enum):
                 SALES = "sales"
                 SUPPORT = "support"
@@ -100,11 +105,4 @@ def router(transcript: str) -> Department:
                     instructions="Select the best department for the customer request",
                 )
 
-            department = router("I need to update my payment method")
-            assert department == Department.BILLING
-
-            department = router("Do you price match?")
-            assert department == Department.SALES
-
-            department = router("*angry noises*")
-            assert department == Department.SUPPORT
+            assert router(user_input).value == expected_selection