gemini_cli_session_test.go 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. //go:build unit
  2. package handler
  3. import (
  4. "crypto/sha256"
  5. "encoding/hex"
  6. "net/http/httptest"
  7. "testing"
  8. "github.com/gin-gonic/gin"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func TestExtractGeminiCLISessionHash(t *testing.T) {
  12. tests := []struct {
  13. name string
  14. body string
  15. privilegedUserID string
  16. wantEmpty bool
  17. wantHash string
  18. }{
  19. {
  20. name: "with privileged-user-id and tmp dir",
  21. body: `{"contents":[{"parts":[{"text":"The project's temporary directory is: /Users/ianshaw/.gemini/tmp/f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740"}]}]}`,
  22. privilegedUserID: "90785f52-8bbe-4b17-b111-a1ddea1636c3",
  23. wantEmpty: false,
  24. wantHash: func() string {
  25. combined := "90785f52-8bbe-4b17-b111-a1ddea1636c3:f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740"
  26. hash := sha256.Sum256([]byte(combined))
  27. return hex.EncodeToString(hash[:])
  28. }(),
  29. },
  30. {
  31. name: "without privileged-user-id but with tmp dir",
  32. body: `{"contents":[{"parts":[{"text":"The project's temporary directory is: /Users/ianshaw/.gemini/tmp/f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740"}]}]}`,
  33. privilegedUserID: "",
  34. wantEmpty: false,
  35. wantHash: "f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740",
  36. },
  37. {
  38. name: "without tmp dir",
  39. body: `{"contents":[{"parts":[{"text":"Hello world"}]}]}`,
  40. privilegedUserID: "90785f52-8bbe-4b17-b111-a1ddea1636c3",
  41. wantEmpty: true,
  42. },
  43. {
  44. name: "empty body",
  45. body: "",
  46. privilegedUserID: "90785f52-8bbe-4b17-b111-a1ddea1636c3",
  47. wantEmpty: true,
  48. },
  49. }
  50. for _, tt := range tests {
  51. t.Run(tt.name, func(t *testing.T) {
  52. // 创建测试上下文
  53. w := httptest.NewRecorder()
  54. c, _ := gin.CreateTestContext(w)
  55. c.Request = httptest.NewRequest("POST", "/test", nil)
  56. if tt.privilegedUserID != "" {
  57. c.Request.Header.Set("x-gemini-api-privileged-user-id", tt.privilegedUserID)
  58. }
  59. // 调用函数
  60. result := extractGeminiCLISessionHash(c, []byte(tt.body))
  61. // 验证结果
  62. if tt.wantEmpty {
  63. require.Empty(t, result, "expected empty session hash")
  64. } else {
  65. require.NotEmpty(t, result, "expected non-empty session hash")
  66. require.Equal(t, tt.wantHash, result, "session hash mismatch")
  67. }
  68. })
  69. }
  70. }
  71. func TestGeminiCLITmpDirRegex(t *testing.T) {
  72. tests := []struct {
  73. name string
  74. input string
  75. wantMatch bool
  76. wantHash string
  77. }{
  78. {
  79. name: "valid tmp dir path",
  80. input: "/Users/ianshaw/.gemini/tmp/f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740",
  81. wantMatch: true,
  82. wantHash: "f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740",
  83. },
  84. {
  85. name: "valid tmp dir path in text",
  86. input: "The project's temporary directory is: /Users/ianshaw/.gemini/tmp/f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740\nOther text",
  87. wantMatch: true,
  88. wantHash: "f7851b009ed314d1baee62e83115f486160283f4a55a582d89fdac8b9fe3b740",
  89. },
  90. {
  91. name: "invalid hash length",
  92. input: "/Users/ianshaw/.gemini/tmp/abc123",
  93. wantMatch: false,
  94. },
  95. {
  96. name: "no tmp dir",
  97. input: "Hello world",
  98. wantMatch: false,
  99. },
  100. }
  101. for _, tt := range tests {
  102. t.Run(tt.name, func(t *testing.T) {
  103. match := geminiCLITmpDirRegex.FindStringSubmatch(tt.input)
  104. if tt.wantMatch {
  105. require.NotNil(t, match, "expected regex to match")
  106. require.Len(t, match, 2, "expected 2 capture groups")
  107. require.Equal(t, tt.wantHash, match[1], "hash mismatch")
  108. } else {
  109. require.Nil(t, match, "expected regex not to match")
  110. }
  111. })
  112. }
  113. }
  114. func TestSafeShortPrefix(t *testing.T) {
  115. tests := []struct {
  116. name string
  117. input string
  118. n int
  119. want string
  120. }{
  121. {name: "空字符串", input: "", n: 8, want: ""},
  122. {name: "长度小于截断值", input: "abc", n: 8, want: "abc"},
  123. {name: "长度等于截断值", input: "12345678", n: 8, want: "12345678"},
  124. {name: "长度大于截断值", input: "1234567890", n: 8, want: "12345678"},
  125. {name: "截断值为0", input: "123456", n: 0, want: "123456"},
  126. }
  127. for _, tt := range tests {
  128. t.Run(tt.name, func(t *testing.T) {
  129. require.Equal(t, tt.want, safeShortPrefix(tt.input, tt.n))
  130. })
  131. }
  132. }