material_manage_system/internal/handler/organization_context_test.go

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")
	}
}