openai_quota_platform_contract_test.go 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. package handler
  2. import (
  3. "go/ast"
  4. "go/parser"
  5. "go/token"
  6. "path/filepath"
  7. "testing"
  8. "github.com/stretchr/testify/require"
  9. )
  10. func TestOpenAIRecordUsageInputsCarryQuotaPlatform(t *testing.T) {
  11. files := []string{
  12. "openai_gateway_handler.go",
  13. "openai_chat_completions.go",
  14. "openai_embeddings.go",
  15. "openai_images.go",
  16. }
  17. for _, name := range files {
  18. t.Run(name, func(t *testing.T) {
  19. fset := token.NewFileSet()
  20. file, err := parser.ParseFile(fset, filepath.Join(".", name), nil, 0)
  21. require.NoError(t, err)
  22. var missing []token.Position
  23. ast.Inspect(file, func(node ast.Node) bool {
  24. literal, ok := node.(*ast.CompositeLit)
  25. if !ok || !isOpenAIRecordUsageInputLiteral(literal.Type) {
  26. return true
  27. }
  28. if !compositeLiteralHasKey(literal, "QuotaPlatform") {
  29. missing = append(missing, fset.Position(literal.Lbrace))
  30. }
  31. return true
  32. })
  33. require.Empty(t, missing, "OpenAI usage post-billing must receive request-time QuotaPlatform")
  34. })
  35. }
  36. }
  37. func isOpenAIRecordUsageInputLiteral(expr ast.Expr) bool {
  38. selector, ok := expr.(*ast.SelectorExpr)
  39. if !ok {
  40. return false
  41. }
  42. pkg, ok := selector.X.(*ast.Ident)
  43. return ok && pkg.Name == "service" && selector.Sel.Name == "OpenAIRecordUsageInput"
  44. }
  45. func compositeLiteralHasKey(literal *ast.CompositeLit, key string) bool {
  46. for _, elt := range literal.Elts {
  47. pair, ok := elt.(*ast.KeyValueExpr)
  48. if !ok {
  49. continue
  50. }
  51. ident, ok := pair.Key.(*ast.Ident)
  52. if ok && ident.Name == key {
  53. return true
  54. }
  55. }
  56. return false
  57. }