FilesExpand file tree

main

/

test_mode_handlers_e2e.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

188 lines (159 loc) · 6.61 KB

main

/

test_mode_handlers_e2e.py

Copy path

Top

File metadata and controls

• Code
• Blame

188 lines (159 loc) · 6.61 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

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

"""E2E tests for mode handlers."""

from __future__ import annotations

import asyncio

import pytest

from copilot.rpc import ModeSetRequest

from copilot.session import PermissionHandler

from copilot.session_events import (

AutoModeSwitchCompletedData,

AutoModeSwitchRequestedData,

AutoModeSwitchResponse,

ExitPlanModeAction,

ExitPlanModeCompletedData,

ExitPlanModeRequestedData,

SessionIdleData,

SessionMode,

SessionModelChangeData,

)

from .testharness import E2ETestContext, wait_for_event

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

MODE_HANDLER_TOKEN = "mode-handler-token"

PLAN_SUMMARY = "Greeting file implementation plan"

PLAN_PROMPT = (

"Create a brief implementation plan for adding a greeting.txt file, then request "

"approval with exit_plan_mode."

)

AUTO_MODE_PROMPT = "Explain that auto mode recovered from a rate limit in one short sentence."

@pytest.fixture(scope="module")

async def mode_ctx(ctx: E2ETestContext):

"""Configure per-token user responses for mode-handler tests."""

proxy_url = ctx.proxy_url

ctx.add_runtime_env("COPILOT_DEBUG_GITHUB_API_URL", proxy_url)

await ctx.set_copilot_user_by_token(

MODE_HANDLER_TOKEN,

{

"login": "mode-handler-user",

"copilot_plan": "individual_pro",

"endpoints": {

"api": proxy_url,

"telemetry": "https://sdk-hydro-collector.example.com",

},

"analytics_tracking_id": "mode-handler-tracking-id",

},

)

return ctx

class TestModeHandlers:

async def test_should_invoke_exit_plan_mode_handler_when_model_uses_tool(

self, mode_ctx: E2ETestContext

):

exit_plan_mode_requests = []

async def on_exit_plan_mode_request(request, invocation):

exit_plan_mode_requests.append(request)

assert invocation["session_id"] == session.session_id

return {

"approved": True,

"selectedAction": "interactive",

"feedback": "Approved by the Python E2E test",

}

session = await mode_ctx.client.create_session(

github_token=MODE_HANDLER_TOKEN,

on_permission_request=PermissionHandler.approve_all,

on_exit_plan_mode_request=on_exit_plan_mode_request,

)

requested_event = wait_for_event(

session,

lambda event: (

isinstance(event.data, ExitPlanModeRequestedData)

and event.data.summary == PLAN_SUMMARY

),

)

completed_event = wait_for_event(

session,

lambda event: (

isinstance(event.data, ExitPlanModeCompletedData)

and event.data.approved is True

and event.data.selected_action == ExitPlanModeAction.INTERACTIVE

),

)

try:

await session.rpc.mode.set(ModeSetRequest(mode=SessionMode.PLAN))

response = await session.send_and_wait(PLAN_PROMPT)

assert len(exit_plan_mode_requests) == 1

request = exit_plan_mode_requests[0]

assert request["summary"] == PLAN_SUMMARY

assert request["actions"] == ["autopilot", "interactive", "exit_only"]

assert request["recommendedAction"] == "interactive"

assert request.get("planContent") is not None

requested = await requested_event

assert requested.data.summary == PLAN_SUMMARY

completed = await completed_event

assert completed.data.approved is True

assert completed.data.selected_action == ExitPlanModeAction.INTERACTIVE

assert completed.data.feedback == "Approved by the Python E2E test"

assert response is not None

finally:

requested_event.cancel()

completed_event.cancel()

await asyncio.gather(requested_event, completed_event, return_exceptions=True)

await session.disconnect()

async def test_should_invoke_auto_mode_switch_handler_when_rate_limited(

self, mode_ctx: E2ETestContext

):

auto_mode_switch_requests = []

async def on_auto_mode_switch_request(request, invocation):

auto_mode_switch_requests.append(request)

assert invocation["session_id"] == session.session_id

return "yes"

session = await mode_ctx.client.create_session(

github_token=MODE_HANDLER_TOKEN,

on_permission_request=PermissionHandler.approve_all,

on_auto_mode_switch_request=on_auto_mode_switch_request,

)

requested_event = wait_for_event(

session,

lambda event: (

isinstance(event.data, AutoModeSwitchRequestedData)

and event.data.error_code == "user_weekly_rate_limited"

and event.data.retry_after_seconds == 1

),

)

completed_event = wait_for_event(

session,

lambda event: (

isinstance(event.data, AutoModeSwitchCompletedData)

and event.data.response == AutoModeSwitchResponse.YES

),

)

model_change_event = wait_for_event(

session,

lambda event: (

isinstance(event.data, SessionModelChangeData)

and event.data.cause == "rate_limit_auto_switch"

),

)

idle_event = wait_for_event(

session,

lambda event: isinstance(event.data, SessionIdleData),

)

try:

message_id = await session.send(AUTO_MODE_PROMPT)

assert message_id

requested = await requested_event

assert requested.data.error_code == "user_weekly_rate_limited"

assert requested.data.retry_after_seconds == 1

completed = await completed_event

assert completed.data.response == AutoModeSwitchResponse.YES

model_change = await model_change_event

assert model_change.data.cause == "rate_limit_auto_switch"

idle = await idle_event

assert isinstance(idle.data, SessionIdleData)

assert len(auto_mode_switch_requests) == 1

request = auto_mode_switch_requests[0]

assert request["errorCode"] == "user_weekly_rate_limited"

assert request["retryAfterSeconds"] == 1

finally:

waiters = [requested_event, completed_event, model_change_event, idle_event]

for waiter in waiters:

waiter.cancel()

await asyncio.gather(*waiters, return_exceptions=True)

await session.disconnect()