update tests

microsoft · Apr 1, 2024 · 4c748d4 · 4c748d4
1 parent 84cc10f
commit 4c748d4
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/src/promptflow/tests/sdk_pfs_test/e2etests/test_connection_apis.py b/src/promptflow/tests/sdk_pfs_test/e2etests/test_connection_apis.py
@@ -11,6 +11,7 @@
 
 from promptflow import PFClient
 from promptflow._sdk.entities import CustomConnection
+from promptflow._sdk._version import VERSION
 from promptflow.recording.record_mode import is_replay
 
 from ..utils import PFSOperations, check_activity_end_telemetry
@@ -34,7 +35,7 @@ def test_list_connections(self, pf_client: PFClient, pfs_op: PFSOperations) -> N
 
     def test_list_connections_with_different_user_agent(self, pf_client: PFClient, pfs_op: PFSOperations) -> None:
         create_custom_connection(pf_client)
-        base_user_agent = ["local_pfs/0.0.1"]
+        base_user_agent = [f"local_pfs/{VERSION}"]
         for _, extra_user_agent in enumerate(
             [
                 ["another_test_user_agent/0.0.1"],

diff --git a/src/promptflow/tests/sdk_pfs_test/e2etests/test_telemetry_apis.py b/src/promptflow/tests/sdk_pfs_test/e2etests/test_telemetry_apis.py
@@ -4,6 +4,8 @@
 
 import pytest
 
+from promptflow._sdk._version import VERSION
+
 from ..utils import PFSOperations, check_activity_end_telemetry
 
 
@@ -34,7 +36,7 @@ def test_post_telemetry(self, pfs_op: PFSOperations) -> None:
         with check_activity_end_telemetry(
             activity_name="pf.flow.test",
             activity_type="InternalCall",
-            user_agent=f"{user_agent} local_pfs/0.0.1",
+            user_agent=f"{user_agent} local_pfs/{VERSION}",
             request_id=request_id,
         ):
             response = pfs_op.create_telemetry(

diff --git a/src/promptflow/tests/sdk_pfs_test/utils.py b/src/promptflow/tests/sdk_pfs_test/utils.py
@@ -11,6 +11,7 @@
 from flask.testing import FlaskClient
 
 from promptflow._sdk._service.utils.utils import encrypt_flow_path
+from promptflow._sdk._version import VERSION
 
 
 @contextlib.contextmanager
@@ -33,7 +34,7 @@ def check_activity_end_telemetry(
             "first_call": True,
             "activity_type": "PublicApi",
             "completion_status": "Success",
-            "user_agent": [f"Werkzeug/{werkzeug.__version__}", "local_pfs/0.0.1"],
+            "user_agent": [f"Werkzeug/{werkzeug.__version__}", f"local_pfs/{VERSION}"],
         }
         for i, expected_activity in enumerate(expected_activities):
             temp = default_expected_call.copy()