llm-intelligence/scripts/import_vertex_pricing_test.go

//go:build llm_script

package main

import (
	"bytes"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestParseVertexPricingCatalogBuildsRecords(t *testing.T) {
	raw, err := os.ReadFile(filepath.Join("testdata", "vertex_pricing_sample.html"))
	if err != nil {
		t.Fatalf("读取 fixture 失败: %v", err)
	}

	records, err := parseVertexPricingCatalog(string(raw))
	if err != nil {
		t.Fatalf("parseVertexPricingCatalog 返回错误: %v", err)
	}
	if len(records) != 4 {
		t.Fatalf("期望 4 条 Vertex 价格记录，实际 %d", len(records))
	}
	if records[0].ModelName != "Gemini 3.1 Pro Preview" {
		t.Fatalf("首条模型名错误: %q", records[0].ModelName)
	}
	if records[1].InputPrice != 0.5 || records[1].OutputPrice != 3 {
		t.Fatalf("Gemini 3.1 Flash Image 定价错误: %v / %v", records[1].InputPrice, records[1].OutputPrice)
	}
	if records[2].InputPrice != 0.25 || records[2].OutputPrice != 1.5 {
		t.Fatalf("Gemini 3.1 Flash-Lite 定价错误: %v / %v", records[2].InputPrice, records[2].OutputPrice)
	}
}

func TestRunVertexPricingImportDryRunPrintsSummary(t *testing.T) {
	var out bytes.Buffer
	err := runVertexPricingImport(vertexPricingImportConfig{
		URL:     defaultVertexPricingURL,
		Fixture: filepath.Join("testdata", "vertex_pricing_sample.html"),
		DryRun:  true,
	}, nil, &out)
	if err != nil {
		t.Fatalf("runVertexPricingImport 返回错误: %v", err)
	}
	output := out.String()
	for _, want := range []string{
		"source=vertex-pricing-import",
		"models=4",
		"operator=Google Cloud Vertex AI",
		"dry_run=true",
	} {
		if !strings.Contains(output, want) {
			t.Fatalf("输出缺少 %q，实际: %q", want, output)
		}
	}
}

func TestParseVertexPricingCatalogAcceptsGenericStandardTableMarkup(t *testing.T) {
	raw := `
<h2>Gemini 2.5</h2>
<section>
  <h4>Standard</h4>
  <table>
    <tbody>
      <tr>
        <th>Model</th>
        <th>Type</th>
        <th>Price</th>
      </tr>
      <tr>
        <td rowspan="3">Gemini 2.5 Flash</td>
      </tr>
      <tr>
        <td>Input (text, image, video)</td>
        <td>$0.30</td>
      </tr>
      <tr>
        <td>Text output</td>
        <td>$2.50</td>
      </tr>
    </tbody>
  </table>
</section>
`

	records, err := parseVertexPricingCatalog(raw)
	if err != nil {
		t.Fatalf("parseVertexPricingCatalog 返回错误: %v", err)
	}
	if len(records) != 1 {
		t.Fatalf("期望 1 条 Vertex 价格记录，实际 %d", len(records))
	}
	if records[0].ModelName != "Gemini 2.5 Flash" {
		t.Fatalf("模型名错误: %q", records[0].ModelName)
	}
	if records[0].InputPrice != 0.3 || records[0].OutputPrice != 2.5 {
		t.Fatalf("价格解析错误: %v / %v", records[0].InputPrice, records[0].OutputPrice)
	}
}

func TestParseVertexPricingCatalogFallsBackToStandardTextBlocks(t *testing.T) {
	raw := `
<div>### Standard</div>
<div>Model Type Price (/1M tokens) <= 200K input tokens Price (/1M tokens) > 200K input tokens</div>
<div>Gemini 2.5</div>
<div>Flash</div>
<div>Input (text, image, video) $0.54 $0.54 $0.05 $0.05</div>
<div>Audio Input $1.80 $1.80 $0.18 $0.18</div>
<div>Text output (response and reasoning) $4.50 $4.50 N/A N/A</div>
<div>### Flex/Batch</div>
`

	records, err := parseVertexPricingCatalog(raw)
	if err != nil {
		t.Fatalf("parseVertexPricingCatalog 返回错误: %v", err)
	}
	if len(records) != 1 {
		t.Fatalf("期望 1 条 Vertex 价格记录，实际 %d", len(records))
	}
	if records[0].ModelName != "Gemini 2.5 Flash" {
		t.Fatalf("模型名错误: %q", records[0].ModelName)
	}
	if records[0].InputPrice != 0.54 || records[0].OutputPrice != 4.5 {
		t.Fatalf("价格解析错误: %v / %v", records[0].InputPrice, records[0].OutputPrice)
	}
}

func TestParseVertexPricingCatalogSupportsChineseStandardTable(t *testing.T) {
	raw := `
<h3>Gemini 3</h3>
<section>
  <h3 id="standard">标准</h3>
  <table class="style0">
    <tbody>
      <tr>
        <th>模型</th>
        <th>类型</th>
        <th>价格</th>
      </tr>
      <tr>
        <td rowspan="3">Gemini 3 Pro 预览版</td>
      </tr>
      <tr>
        <td>输入（文本、图片、视频、音频）</td>
        <td>$2</td>
      </tr>
      <tr>
        <td>文本输出（回答和推理）</td>
        <td>$12</td>
      </tr>
    </tbody>
  </table>
</section>
`

	records, err := parseVertexPricingCatalog(raw)
	if err != nil {
		t.Fatalf("parseVertexPricingCatalog 返回错误: %v", err)
	}
	if len(records) != 1 {
		t.Fatalf("期望 1 条记录，实际 %d", len(records))
	}
	if records[0].ModelName != "Gemini 3 Pro 预览版" {
		t.Fatalf("模型名错误: %q", records[0].ModelName)
	}
	if records[0].InputPrice != 2 || records[0].OutputPrice != 12 {
		t.Fatalf("价格解析错误: %v / %v", records[0].InputPrice, records[0].OutputPrice)
	}
}