material_manage_system/internal/middleware/organization_test.go

package middleware

import (
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"

	"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 organizationMock(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 TestOrganizationMiddlewareHeadersAndAccess(t *testing.T) {
	gin.SetMode(gin.TestMode)
	for _, tc := range []struct {
		name, header string
		allowed      bool
		status       int
	}{
		{"missing", "", false, 400}, {"zero", "0", false, 400}, {"negative", "-1", false, 400},
		{"signed", "+1", false, 400}, {"space", " 1", false, 400}, {"fraction", "1.2", false, 400},
		{"text", "abc", false, 400}, {"overflow", "9223372036854775808", false, 400},
		{"ungranted", "2", false, 403}, {"nonexistent", "999", false, 403}, {"inactive", "3", false, 403},
		{"allowed", "2", true, 204},
	} {
		t.Run(tc.name, func(t *testing.T) {
			mock := organizationMock(t)
			if tc.status != 400 {
				mock.ExpectQuery("SELECT EXISTS").WithArgs(sqlmock.AnyArg(), int64(7), int64(7)).WillReturnRows(sqlmock.NewRows([]string{"allowed"}).AddRow(tc.allowed))
			}
			router := gin.New()
			router.GET("/", func(c *gin.Context) { c.Set(UserIDKey, int64(7)) }, OrganizationMiddleware(), func(c *gin.Context) {
				if GetCurrentOrganizationID(c) != 2 {
					t.Fatal("trusted organization missing")
				}
				c.Request.Header.Set("X-Organization-ID", "999")
				if GetCurrentOrganizationID(c) != 2 {
					t.Fatal("header mutated trusted context")
				}
				c.Status(204)
			})
			req := httptest.NewRequest("GET", "/", nil)
			req.Header.Set("X-Organization-ID", tc.header)
			w := httptest.NewRecorder()
			router.ServeHTTP(w, req)
			if w.Code != tc.status {
				t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
			}
			if err := mock.ExpectationsWereMet(); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestOrganizationMiddlewareRechecksRevocationOnNextRequest(t *testing.T) {
	for _, name := range []string{"grant revoked", "company disabled", "super admin revoked"} {
		t.Run(name, func(t *testing.T) {
			mock := organizationMock(t)
			token, err := service.NewAuthService().GenerateToken(&model.User{ID: 7, Roles: []model.Role{{Code: "super_admin", Enabled: true}}})
			if err != nil {
				t.Fatal(err)
			}
			router := gin.New()
			router.GET("/", AuthMiddleware(), OrganizationMiddleware(), func(c *gin.Context) { c.Status(204) })
			for i, allowed := range []bool{true, false} {
				role := "warehouse_admin"
				if name == "super admin revoked" && i == 0 {
					role = "super_admin"
				}
				mock.ExpectQuery("SELECT u.id, r.code FROM users u").WithArgs(int64(7)).WillReturnRows(sqlmock.NewRows([]string{"id", "code"}).AddRow(7, role))
				mock.ExpectQuery("SELECT EXISTS").WithArgs(int64(2), int64(7), int64(7)).WillReturnRows(sqlmock.NewRows([]string{"allowed"}).AddRow(allowed))
				req := httptest.NewRequest("GET", "/", nil)
				req.Header.Set("Authorization", "Bearer "+token)
				req.Header.Set("X-Organization-ID", "2")
				w := httptest.NewRecorder()
				router.ServeHTTP(w, req)
				want := 204
				if !allowed {
					want = 403
				}
				if w.Code != want {
					t.Fatalf("request %d status=%d", i, w.Code)
				}
			}
			if err := mock.ExpectationsWereMet(); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestAuthMiddlewareUsesCurrentRolesAndFailsClosed(t *testing.T) {
	for _, tc := range []struct {
		name     string
		role     *string
		queryErr error
		missing  bool
		status   int
	}{
		{name: "stale admin", role: stringPointer("warehouse_admin"), status: 403},
		{name: "disabled roles", status: 403}, {name: "deleted user", missing: true, status: 401},
		{name: "database failure", queryErr: errors.New("offline"), status: 500},
		{name: "current admin", role: stringPointer("super_admin"), status: 204},
	} {
		t.Run(tc.name, func(t *testing.T) {
			mock := organizationMock(t)
			query := mock.ExpectQuery("SELECT u.id, r.code FROM users u").WithArgs(int64(7))
			rows := sqlmock.NewRows([]string{"id", "code"})
			if !tc.missing {
				rows.AddRow(7, tc.role)
			}
			if tc.queryErr != nil {
				query.WillReturnError(tc.queryErr)
			} else {
				query.WillReturnRows(rows)
			}
			token, _ := service.NewAuthService().GenerateToken(&model.User{ID: 7, Roles: []model.Role{{Code: "super_admin", Enabled: true}}})
			router := gin.New()
			router.GET("/", AuthMiddleware(), AdminOnly(), func(c *gin.Context) { c.Status(http.StatusNoContent) })
			req := httptest.NewRequest("GET", "/", nil)
			req.Header.Set("Authorization", "Bearer "+token)
			w := httptest.NewRecorder()
			router.ServeHTTP(w, req)
			if w.Code != tc.status {
				t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
			}
			if err := mock.ExpectationsWereMet(); err != nil {
				t.Fatal(err)
			}
		})
	}
}
func stringPointer(s string) *string { return &s }