package handler import ( "encoding/json" "errors" "net/http/httptest" "strings" "testing" "time" "github.com/DATA-DOG/go-sqlmock" "github.com/gin-gonic/gin" "github.com/wmms/wmms/internal/model" "github.com/wmms/wmms/internal/service" "github.com/wmms/wmms/pkg/database" ) func contextMock(t *testing.T) sqlmock.Sqlmock { t.Helper() db, mock, err := sqlmock.New() if err != nil { t.Fatal(err) } previous := database.DB database.DB = db t.Cleanup(func() { database.DB = previous; db.Close() }) return mock } func contextToken(t *testing.T) string { t.Helper() token, err := service.NewAuthService().GenerateToken(&model.User{ID: 7, Username: "operator"}) if err != nil { t.Fatal(err) } return token } func expectContextPrincipal(mock sqlmock.Sqlmock, role string) { mock.ExpectQuery("SELECT u.id, r.code FROM users u").WithArgs(int64(7)).WillReturnRows(sqlmock.NewRows([]string{"id", "code"}).AddRow(7, role)) } func contextAccessRows(allowed bool) *sqlmock.Rows { rows := sqlmock.NewRows([]string{"id", "code", "name", "status", "default_timezone", "created_at", "updated_at", "is_default"}) if allowed { now := time.Now() rows.AddRow(2, "ORG2", "公司二", "active", "Asia/Shanghai", now, now, true) } return rows } func expectContextAccess(mock sqlmock.Sqlmock, allowed bool) { mock.ExpectQuery("SELECT o.id,o.code,o.name,o.status").WithArgs(int64(7), int64(7)).WillReturnRows(contextAccessRows(allowed)) } func TestOrganizationContextListWithoutOrganizationHeader(t *testing.T) { gin.SetMode(gin.TestMode) mock := contextMock(t) expectContextPrincipal(mock, "warehouse_admin") expectContextAccess(mock, true) router := gin.New() NewHandler().RegisterRoutes(router) req := httptest.NewRequest("GET", "/api/auth/organizations", nil) req.Header.Set("Authorization", "Bearer "+contextToken(t)) w := httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != 200 || !strings.Contains(w.Body.String(), `"default_organization_id":2`) { t.Fatalf("status=%d body=%s", w.Code, w.Body.String()) } if err := mock.ExpectationsWereMet(); err != nil { t.Fatal(err) } } func TestOrganizationContextSwitchValidatesTargetAndRequiresAudit(t *testing.T) { for _, tc := range []struct { name, body, config string allowed, auditFailure bool status int }{ {"no warehouse", `{"organization_id":2,"previous_organization_id":999}`, "absent", true, false, 200}, {"own disabled warehouse", `{"organization_id":2,"previous_organization_id":999}`, "disabled", true, false, 200}, {"own warehouse no location", `{"organization_id":2,"previous_organization_id":999}`, "no location", true, false, 200}, {"audit failure", `{"organization_id":2,"previous_organization_id":999}`, "absent", true, true, 500}, {"ungranted target", `{"organization_id":2,"previous_organization_id":1}`, "absent", false, false, 403}, {"invalid target", `{"organization_id":0}`, "absent", false, false, 400}, {"ambiguous warehouse", `{"organization_id":2}`, "multi warehouse", true, false, 500}, {"ambiguous location", `{"organization_id":2}`, "multi location", true, false, 500}, } { t.Run(tc.name, func(t *testing.T) { mock := contextMock(t) expectContextPrincipal(mock, "warehouse_admin") if tc.status != 400 { expectContextAccess(mock, tc.allowed) } if tc.allowed { rows := sqlmock.NewRows([]string{"id", "organization_id", "code", "name", "status", "location_id", "location_code", "location_name", "location_status"}) switch tc.config { case "disabled": rows.AddRow(8, 2, "OWN", "公司二仓库", "inactive", 9, "OWN-L", "公司二库位", "inactive") case "no location": rows.AddRow(8, 2, "OWN", "公司二仓库", "active", nil, nil, nil, nil) case "multi warehouse": rows.AddRow(8, 2, "A", "仓一", "active", nil, nil, nil, nil).AddRow(9, 2, "B", "仓二", "active", nil, nil, nil, nil) case "multi location": rows.AddRow(8, 2, "A", "仓一", "active", 9, "L1", "位一", "active").AddRow(8, 2, "A", "仓一", "active", 10, "L2", "位二", "active") } mock.ExpectQuery("SELECT w.id, w.organization_id").WithArgs(int64(2)).WillReturnRows(rows) if tc.status == 200 || tc.auditFailure { audit := mock.ExpectExec("INSERT INTO organization_audit_logs").WithArgs(int64(7), int64(2), "organization_context_switched", `{"organization_id":999}`, `{"organization_id":2}`) if tc.auditFailure { audit.WillReturnError(errors.New("audit unavailable")) } else { audit.WillReturnResult(sqlmock.NewResult(10, 1)) } } } router := gin.New() NewHandler().RegisterRoutes(router) req := httptest.NewRequest("POST", "/api/auth/organization-context", strings.NewReader(tc.body)) req.Header.Set("Authorization", "Bearer "+contextToken(t)) req.Header.Set("Content-Type", "application/json") // The target must be validated from the body even if a forged header disagrees. req.Header.Set("X-Organization-ID", "999") w := httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != tc.status { t.Fatalf("status=%d body=%s", w.Code, w.Body.String()) } if tc.status == 200 { var envelope struct { Data service.OrganizationContext `json:"data"` } if err := json.Unmarshal(w.Body.Bytes(), &envelope); err != nil { t.Fatal(err) } if envelope.Data.OrganizationID != 2 || envelope.Data.Organization.ID != 2 { t.Fatal("wrong company") } warehouse := envelope.Data.Warehouse if tc.config == "absent" && warehouse != nil { t.Fatal("must not fall back to a default warehouse") } if tc.config == "disabled" && (warehouse == nil || warehouse.ID != 8 || warehouse.Status != "inactive" || warehouse.Location == nil || warehouse.Location.ID != 9 || warehouse.Location.Status != "inactive") { t.Fatal("wrong disabled summary") } if tc.config == "no location" && (warehouse == nil || warehouse.Location != nil) { t.Fatal("missing location must remain null") } } if err := mock.ExpectationsWereMet(); err != nil { t.Fatal(err) } }) } } func TestOrganizationIDReadsOnlyTrustedContext(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest("GET", "/", nil) c.Request.Header.Set("X-Organization-ID", "999") if id, ok := organizationID(c); ok || id != 0 { t.Fatal("unverified header granted company") } c.Set("organization_id", int64(2)) if id, ok := organizationID(c); !ok || id != 2 { t.Fatal("trusted context ignored") } }