gemini_native_signature_cleaner_test.go 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. package service
  2. import (
  3. "encoding/json"
  4. "testing"
  5. "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
  6. "github.com/stretchr/testify/require"
  7. )
  8. func TestCleanGeminiNativeThoughtSignatures_ReplacesNestedThoughtSignatures(t *testing.T) {
  9. input := []byte(`{
  10. "contents": [
  11. {
  12. "role": "user",
  13. "parts": [{"text": "hello"}]
  14. },
  15. {
  16. "role": "model",
  17. "parts": [
  18. {"text": "thinking", "thought": true, "thoughtSignature": "sig_1"},
  19. {"functionCall": {"name": "toolA", "args": {"k": "v"}}, "thoughtSignature": "sig_2"}
  20. ]
  21. }
  22. ],
  23. "cachedContent": {
  24. "parts": [{"text": "cached", "thoughtSignature": "sig_3"}]
  25. },
  26. "signature": "keep_me"
  27. }`)
  28. cleaned := CleanGeminiNativeThoughtSignatures(input)
  29. var got map[string]any
  30. require.NoError(t, json.Unmarshal(cleaned, &got))
  31. require.NotContains(t, string(cleaned), `"thoughtSignature":"sig_1"`)
  32. require.NotContains(t, string(cleaned), `"thoughtSignature":"sig_2"`)
  33. require.NotContains(t, string(cleaned), `"thoughtSignature":"sig_3"`)
  34. require.Contains(t, string(cleaned), `"thoughtSignature":"`+antigravity.DummyThoughtSignature+`"`)
  35. require.Contains(t, string(cleaned), `"signature":"keep_me"`)
  36. }
  37. func TestCleanGeminiNativeThoughtSignatures_InvalidJSONReturnsOriginal(t *testing.T) {
  38. input := []byte(`{"contents":[invalid-json]}`)
  39. cleaned := CleanGeminiNativeThoughtSignatures(input)
  40. require.Equal(t, input, cleaned)
  41. }
  42. func TestReplaceThoughtSignaturesRecursive_OnlyReplacesTargetField(t *testing.T) {
  43. input := map[string]any{
  44. "thoughtSignature": "sig_root",
  45. "signature": "keep_signature",
  46. "nested": []any{
  47. map[string]any{
  48. "thoughtSignature": "sig_nested",
  49. "signature": "keep_nested_signature",
  50. },
  51. },
  52. }
  53. got, ok := replaceThoughtSignaturesRecursive(input).(map[string]any)
  54. require.True(t, ok)
  55. require.Equal(t, antigravity.DummyThoughtSignature, got["thoughtSignature"])
  56. require.Equal(t, "keep_signature", got["signature"])
  57. nested, ok := got["nested"].([]any)
  58. require.True(t, ok)
  59. nestedMap, ok := nested[0].(map[string]any)
  60. require.True(t, ok)
  61. require.Equal(t, antigravity.DummyThoughtSignature, nestedMap["thoughtSignature"])
  62. require.Equal(t, "keep_nested_signature", nestedMap["signature"])
  63. }