FilesExpand file tree

main

/

test_system_message_transform_e2e.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

123 lines (92 loc) · 3.92 KB

main

/

test_system_message_transform_e2e.py

Copy path

Top

File metadata and controls

• Code
• Blame

123 lines (92 loc) · 3.92 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

"""

Copyright (c) Microsoft Corporation.

Tests for system message transform functionality

"""

import pytest

from copilot.session import PermissionHandler

from .testharness import E2ETestContext

from .testharness.helper import write_file

pytestmark = pytest.mark.asyncio(loop_scope="module")

class TestSystemMessageTransform:

async def test_should_invoke_transform_callbacks_with_section_content(

self, ctx: E2ETestContext

):

"""Test that transform callbacks are invoked with the section content"""

identity_contents = []

tone_contents = []

async def identity_transform(content: str) -> str:

identity_contents.append(content)

return content

async def tone_transform(content: str) -> str:

tone_contents.append(content)

return content

session = await ctx.client.create_session(

system_message={

"mode": "customize",

"sections": {

"identity": {"action": identity_transform},

"tone": {"action": tone_transform},

},

},

on_permission_request=PermissionHandler.approve_all,

)

write_file(ctx.work_dir, "test.txt", "Hello transform!")

await session.send_and_wait("Read the contents of test.txt and tell me what it says")

# Both transform callbacks should have been invoked

assert len(identity_contents) > 0

assert len(tone_contents) > 0

# Callbacks should have received non-empty content

assert all(len(c) > 0 for c in identity_contents)

assert all(len(c) > 0 for c in tone_contents)

await session.disconnect()

async def test_should_apply_transform_modifications_to_section_content(

self, ctx: E2ETestContext

):

"""Test that transform modifications are applied to the section content"""

async def identity_transform(content: str) -> str:

return content + "\nTRANSFORM_MARKER"

session = await ctx.client.create_session(

system_message={

"mode": "customize",

"sections": {

"identity": {"action": identity_transform},

},

},

on_permission_request=PermissionHandler.approve_all,

)

write_file(ctx.work_dir, "hello.txt", "Hello!")

await session.send_and_wait("Read the contents of hello.txt")

# Verify the transform result was actually applied to the system message

traffic = await ctx.get_exchanges()

system_message = _get_system_message(traffic[0])

assert "TRANSFORM_MARKER" in system_message

await session.disconnect()

async def test_should_work_with_static_overrides_and_transforms_together(

self, ctx: E2ETestContext

):

"""Test that static overrides and transforms work together"""

identity_contents = []

async def identity_transform(content: str) -> str:

identity_contents.append(content)

return content

session = await ctx.client.create_session(

system_message={

"mode": "customize",

"sections": {

"safety": {"action": "remove"},

"identity": {"action": identity_transform},

},

},

on_permission_request=PermissionHandler.approve_all,

)

write_file(ctx.work_dir, "combo.txt", "Combo test!")

await session.send_and_wait("Read the contents of combo.txt and tell me what it says")

# The transform callback should have been invoked

assert len(identity_contents) > 0

await session.disconnect()

def _get_system_message(exchange: dict) -> str:

messages = exchange.get("request", {}).get("messages", [])

for msg in messages:

if msg.get("role") == "system":

return msg.get("content", "")

return ""