ymt_v3_api-20260721-143604/ymt_v3_api/example_test.go

package ymt_v3_api

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

// 生成测试用的 RSA 密钥对和 AES 密钥
func generateTestKeys() (appID, privateKeyPEM, publicKeyPEM, aesKey string, err error) {
	appID = "test_app_id"
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return
	}
	privBytes, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return
	}
	privateKeyPEM = string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: privBytes}))

	pubBytes, err := x509.MarshalPKIXPublicKey(&priv.PublicKey)
	if err != nil {
		return
	}
	publicKeyPEM = string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubBytes}))

	aesKey = "1234567890123456" // 16 字节
	return
}

// 启动一个模拟服务器，返回加密的成功响应
func setupMockServer(t *testing.T, client *Client, bizResp interface{}) *httptest.Server {
	t.Helper()
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 验证必要 Header
		if r.Header.Get("Appid") == "" || r.Header.Get("Timestamp") == "" || r.Header.Get("Sign") == "" {
			w.WriteHeader(http.StatusBadRequest)
			json.NewEncoder(w).Encode(apiResponse{Code: 400, Message: "missing header"})
			return
		}
		// 解密请求体，验证 ciphertext 存在
		var reqBody cipherRequest
		if err := json.NewDecoder(r.Body).Decode(&reqBody); err != nil || reqBody.Ciphertext == "" {
			w.WriteHeader(http.StatusBadRequest)
			json.NewEncoder(w).Encode(apiResponse{Code: 400, Message: "invalid body"})
			return
		}
		// 构造加密响应
		ciphertext, err := client.encryptPayload(bizResp)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		resp := apiResponse{
			Code:    200,
			Message: "成功",
			Data:    &cipherData{Ciphertext: ciphertext},
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(resp)
	}))
	return server
}

func TestOrderKey(t *testing.T) {
	appID, privPEM, pubPEM, aesKey, err := generateTestKeys()
	if err != nil {
		t.Fatal(err)
	}
	client, err := NewClient(appID, privPEM, pubPEM, aesKey, WithEncryptType("AES"))
	if err != nil {
		t.Fatal(err)
	}
	expectedResp := &OrderKeyResponse{
		OutBizNo:        "order_001",
		TradeNo:         "7251449503000383488",
		Key:             "aZKdU9BymzR6qGRzJM",
		ValidBeginTime:  "2026-06-22 15:30:00",
		ValidEndTime:    "2026-12-31 23:59:59",
		UsableNum:       1,
		UsageNum:        0,
		Status:          1,
		SettlementPrice: 9.9,
		Account:         "18666666666",
	}
	server := setupMockServer(t, client, expectedResp)
	defer server.Close()
	client.baseURL = server.URL

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	req := &OrderKeyRequest{
		OutBizNo:   "order_001",
		ActivityNo: "ACT20260622001",
		Account:    "18666666666",
		NotifyURL:  "https://notify.example.com/openapi",
	}
	resp, err := client.OrderKey(ctx, req)
	if err != nil {
		t.Fatalf("OrderKey failed: %v", err)
	}
	if resp.TradeNo != expectedResp.TradeNo {
		t.Errorf("expected trade_no %s, got %s", expectedResp.TradeNo, resp.TradeNo)
	}
	if resp.Key != expectedResp.Key {
		t.Errorf("expected key %s, got %s", expectedResp.Key, resp.Key)
	}
}

func TestQueryKey(t *testing.T) {
	appID, privPEM, pubPEM, aesKey, err := generateTestKeys()
	if err != nil {
		t.Fatal(err)
	}
	client, err := NewClient(appID, privPEM, pubPEM, aesKey)
	if err != nil {
		t.Fatal(err)
	}
	expectedResp := &QueryKeyResponse{
		OutBizNo: "order_001",
		TradeNo:  "7251449503000383488",
		Status:   1,
	}
	server := setupMockServer(t, client, expectedResp)
	defer server.Close()
	client.baseURL = server.URL

	ctx := context.Background()
	req := &QueryKeyRequest{OutBizNo: "order_001"}
	resp, err := client.QueryKey(ctx, req)
	if err != nil {
		t.Fatalf("QueryKey failed: %v", err)
	}
	if resp.OutBizNo != expectedResp.OutBizNo {
		t.Errorf("expected out_biz_no %s, got %s", expectedResp.OutBizNo, resp.OutBizNo)
	}
}

func TestDiscardKey(t *testing.T) {
	appID, privPEM, pubPEM, aesKey, err := generateTestKeys()
	if err != nil {
		t.Fatal(err)
	}
	client, err := NewClient(appID, privPEM, pubPEM, aesKey)
	if err != nil {
		t.Fatal(err)
	}
	expectedResp := &DiscardKeyResponse{
		OutBizNo: "order_001",
		TradeNo:  "7251449503000383488",
		Status:   3,
	}
	server := setupMockServer(t, client, expectedResp)
	defer server.Close()
	client.baseURL = server.URL

	ctx := context.Background()
	req := &DiscardKeyRequest{TradeNo: "7251449503000383488"}
	resp, err := client.DiscardKey(ctx, req)
	if err != nil {
		t.Fatalf("DiscardKey failed: %v", err)
	}
	if resp.Status != 3 {
		t.Errorf("expected status 3, got %d", resp.Status)
	}
}

func TestBatchOrderKey(t *testing.T) {
	appID, privPEM, pubPEM, aesKey, err := generateTestKeys()
	if err != nil {
		t.Fatal(err)
	}
	client, err := NewClient(appID, privPEM, pubPEM, aesKey)
	if err != nil {
		t.Fatal(err)
	}
	expectedResp := &BatchOrderResponse{
		OutBizNo: "batch_001",
		TradeNo:  "7251449503000383499",
		Status:   "processing",
	}
	server := setupMockServer(t, client, expectedResp)
	defer server.Close()
	client.baseURL = server.URL

	ctx := context.Background()
	req := &BatchOrderRequest{
		OutBizNo:   "batch_001",
		ActivityNo: "ACT20260622001",
		Number:     100,
	}
	resp, err := client.BatchOrderKey(ctx, req)
	if err != nil {
		t.Fatalf("BatchOrderKey failed: %v", err)
	}
	if resp.Status != "processing" {
		t.Errorf("expected status processing, got %s", resp.Status)
	}
}

func TestBatchQueryKey(t *testing.T) {
	appID, privPEM, pubPEM, aesKey, err := generateTestKeys()
	if err != nil {
		t.Fatal(err)
	}
	client, err := NewClient(appID, privPEM, pubPEM, aesKey)
	if err != nil {
		t.Fatal(err)
	}
	expectedResp := &BatchQueryResponse{
		OutBizNo:    "batch_001",
		TradeNo:     "7251449503000383499",
		Status:      "success",
		DownloadURL: "https://oss.example.com/openapi_7251449503000383499.zip",
		ZipPassword: "123456",
	}
	server := setupMockServer(t, client, expectedResp)
	defer server.Close()
	client.baseURL = server.URL

	ctx := context.Background()
	req := &BatchQueryRequest{TradeNo: "7251449503000383499"}
	resp, err := client.BatchQueryKey(ctx, req)
	if err != nil {
		t.Fatalf("BatchQueryKey failed: %v", err)
	}
	if resp.DownloadURL != expectedResp.DownloadURL {
		t.Errorf("expected download_url %s, got %s", expectedResp.DownloadURL, resp.DownloadURL)
	}
}