代理批量关联租户接口、关于页运行时与资源指标

internal/service/proxy_test.go

@@ -2,6 +2,7 @@ package service
 
 import (
 	"context"
+	"errors"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -227,3 +228,99 @@ func TestProxyGeoProbe(t *testing.T) {
 		t.Fatalf("probe view = %+v, want 空地区带探测时间", view)
 	}
 }
+
+func TestProxySetTenants(t *testing.T) {
+	svc, db := newProxyEnv(t)
+	ctx := context.Background()
+
+	pa, _ := svc.Create(ctx, ProxyInput{Type: "socks5", Host: "10.0.0.1", Port: 1080})
+	pb, _ := svc.Create(ctx, ProxyInput{Type: "socks5", Host: "10.0.0.2", Port: 1080})
+	// t1 已挂 A,t2 挂 B,t3 直连
+	t1 := model.OciConfig{Alias: "t1", ProxyID: &pa.ID}
+	t2 := model.OciConfig{Alias: "t2", ProxyID: &pb.ID}
+	t3 := model.OciConfig{Alias: "t3"}
+	for _, cfg := range []*model.OciConfig{&t1, &t2, &t3} {
+		if err := db.Create(cfg).Error; err != nil {
+			t.Fatalf("seed config: %v", err)
+		}
+	}
+
+	proxyOf := func(id uint) *uint {
+		var row model.OciConfig
+		if err := db.First(&row, id).Error; err != nil {
+			t.Fatalf("load config %d: %v", id, err)
+		}
+		return row.ProxyID
+	}
+
+	cases := []struct {
+		name    string
+		proxyID uint
+		cfgIDs  []uint
+		wantN   int64
+		wantErr error
+	}{
+		{name: "改挂并新增", proxyID: pa.ID, cfgIDs: []uint{t2.ID, t3.ID, t3.ID}, wantN: 2},
+		{name: "清空解除全部", proxyID: pa.ID, cfgIDs: nil, wantN: 0},
+		{name: "无效租户拒绝", proxyID: pa.ID, cfgIDs: []uint{9999}, wantErr: ErrProxyInvalid},
+		{name: "代理不存在", proxyID: 9999, cfgIDs: nil, wantErr: gorm.ErrRecordNotFound},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			n, err := svc.SetTenants(ctx, tc.proxyID, tc.cfgIDs)
+			if tc.wantErr != nil {
+				if !errors.Is(err, tc.wantErr) {
+					t.Fatalf("err = %v, want %v", err, tc.wantErr)
+				}
+				return
+			}
+			if err != nil || n != tc.wantN {
+				t.Fatalf("SetTenants = %d, %v; want %d, nil", n, err, tc.wantN)
+			}
+		})
+	}
+
+	// 终态:清空后三租户全部直连(t1 在第一步已被解除,t2/t3 在第二步解除)
+	for _, id := range []uint{t1.ID, t2.ID, t3.ID} {
+		if got := proxyOf(id); got != nil {
+			t.Fatalf("config %d proxyID = %v, want nil", id, *got)
+		}
+	}
+}
+
+func TestProxySetTenantsRebind(t *testing.T) {
+	svc, db := newProxyEnv(t)
+	ctx := context.Background()
+
+	pa, _ := svc.Create(ctx, ProxyInput{Type: "socks5", Host: "10.0.1.1", Port: 1080})
+	pb, _ := svc.Create(ctx, ProxyInput{Type: "socks5", Host: "10.0.1.2", Port: 1080})
+	cfg := model.OciConfig{Alias: "steal", ProxyID: &pa.ID}
+	if err := db.Create(&cfg).Error; err != nil {
+		t.Fatalf("seed config: %v", err)
+	}
+
+	if _, err := svc.SetTenants(ctx, pb.ID, []uint{cfg.ID}); err != nil {
+		t.Fatalf("SetTenants: %v", err)
+	}
+	var row model.OciConfig
+	if err := db.First(&row, cfg.ID).Error; err != nil {
+		t.Fatalf("load config: %v", err)
+	}
+	if row.ProxyID == nil || *row.ProxyID != pb.ID {
+		t.Fatalf("proxyID = %v, want %d(改挂到 B)", row.ProxyID, pb.ID)
+	}
+	// A 侧计数应归零,B 侧为 1
+	views, err := svc.List(ctx)
+	if err != nil {
+		t.Fatalf("List: %v", err)
+	}
+	for _, v := range views {
+		want := int64(0)
+		if v.ID == pb.ID {
+			want = 1
+		}
+		if v.UsedBy != want {
+			t.Fatalf("proxy %d UsedBy = %d, want %d", v.ID, v.UsedBy, want)
+		}
+	}
+}