| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465 |
- package handler
- import (
- "go/ast"
- "go/parser"
- "go/token"
- "path/filepath"
- "testing"
- "github.com/stretchr/testify/require"
- )
- func TestOpenAIRecordUsageInputsCarryQuotaPlatform(t *testing.T) {
- files := []string{
- "openai_gateway_handler.go",
- "openai_chat_completions.go",
- "openai_embeddings.go",
- "openai_images.go",
- }
- for _, name := range files {
- t.Run(name, func(t *testing.T) {
- fset := token.NewFileSet()
- file, err := parser.ParseFile(fset, filepath.Join(".", name), nil, 0)
- require.NoError(t, err)
- var missing []token.Position
- ast.Inspect(file, func(node ast.Node) bool {
- literal, ok := node.(*ast.CompositeLit)
- if !ok || !isOpenAIRecordUsageInputLiteral(literal.Type) {
- return true
- }
- if !compositeLiteralHasKey(literal, "QuotaPlatform") {
- missing = append(missing, fset.Position(literal.Lbrace))
- }
- return true
- })
- require.Empty(t, missing, "OpenAI usage post-billing must receive request-time QuotaPlatform")
- })
- }
- }
- func isOpenAIRecordUsageInputLiteral(expr ast.Expr) bool {
- selector, ok := expr.(*ast.SelectorExpr)
- if !ok {
- return false
- }
- pkg, ok := selector.X.(*ast.Ident)
- return ok && pkg.Name == "service" && selector.Sel.Name == "OpenAIRecordUsageInput"
- }
- func compositeLiteralHasKey(literal *ast.CompositeLit, key string) bool {
- for _, elt := range literal.Elts {
- pair, ok := elt.(*ast.KeyValueExpr)
- if !ok {
- continue
- }
- ident, ok := pair.Key.(*ast.Ident)
- if ok && ident.Name == key {
- return true
- }
- }
- return false
- }
|