OpenFlare/atsf_server/router/update_test.go

package router_test

import (
	"atsflare/common"
	"atsflare/router"
	"atsflare/service"
	"bytes"
	"encoding/json"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/gin-contrib/sessions"
	"github.com/gin-contrib/sessions/cookie"
	"github.com/gin-gonic/gin"
)

type roundTripFunc func(req *http.Request) (*http.Response, error)

func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return f(req)
}

func TestLatestReleaseProxy(t *testing.T) {
	gin.SetMode(gin.TestMode)
	common.RedisEnabled = false
	setupTestDB(t)

	originalClient := service.UpdateHTTPClientForTest()
	service.SetUpdateHTTPClientForTest(&http.Client{
		Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
			if req.URL.String() != "https://api.github.com/repos/Rain-kl/ATSFlare/releases/latest" {
				t.Fatalf("unexpected request url: %s", req.URL.String())
			}
			if req.Header.Get("Accept") != "application/vnd.github+json" {
				t.Fatalf("unexpected accept header: %s", req.Header.Get("Accept"))
			}
			if req.Header.Get("User-Agent") != "ATSFlare-Server" {
				t.Fatalf("unexpected user-agent header: %s", req.Header.Get("User-Agent"))
			}
			return &http.Response{
				StatusCode: http.StatusOK,
				Header:     make(http.Header),
				Body: io.NopCloser(strings.NewReader(`{
					"tag_name":"v1.2.3",
					"body":"release notes",
					"html_url":"https://github.com/Rain-kl/ATSFlare/releases/tag/v1.2.3",
					"published_at":"2026-03-11T00:00:00Z"
				}`)),
			}, nil
		}),
	})
	t.Cleanup(func() {
		service.SetUpdateHTTPClientForTest(originalClient)
	})

	engine := gin.New()
	engine.Use(sessions.Sessions("session", cookie.NewStore([]byte("test-secret"))))
	router.SetApiRouter(engine)

	loginBody, err := json.Marshal(map[string]string{
		"username": "root",
		"password": "123456",
	})
	if err != nil {
		t.Fatalf("failed to marshal login body: %v", err)
	}
	loginReq := httptest.NewRequest(http.MethodPost, "/api/user/login", bytes.NewReader(loginBody))
	loginReq.Header.Set("Content-Type", "application/json")
	loginRecorder := httptest.NewRecorder()
	engine.ServeHTTP(loginRecorder, loginReq)
	if loginRecorder.Code != http.StatusOK {
		t.Fatalf("unexpected login status code: %d", loginRecorder.Code)
	}
	loginResult := loginRecorder.Result()
	defer loginResult.Body.Close()

	req := httptest.NewRequest(http.MethodGet, "/api/update/latest-release", nil)
	for _, cookieValue := range loginResult.Cookies() {
		req.AddCookie(cookieValue)
	}

	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)
	if recorder.Code != http.StatusOK {
		t.Fatalf("unexpected status code: %d", recorder.Code)
	}

	var resp apiResponse
	if err := json.Unmarshal(recorder.Body.Bytes(), &resp); err != nil {
		t.Fatalf("failed to decode response: %v", err)
	}
	if !resp.Success {
		t.Fatalf("expected success response, got message: %s", resp.Message)
	}

	var data map[string]any
	if err := json.Unmarshal(resp.Data, &data); err != nil {
		t.Fatalf("failed to decode response data: %v", err)
	}
	if data["tag_name"] != "v1.2.3" {
		t.Fatalf("unexpected tag_name: %#v", data["tag_name"])
	}
	if data["current_version"] != common.Version {
		t.Fatalf("unexpected current_version: %#v", data["current_version"])
	}
}

func loginRootAndBuildEngine(t *testing.T) (*gin.Engine, []*http.Cookie) {
	t.Helper()
	gin.SetMode(gin.TestMode)
	common.RedisEnabled = false
	setupTestDB(t)

	engine := gin.New()
	engine.Use(sessions.Sessions("session", cookie.NewStore([]byte("test-secret"))))
	router.SetApiRouter(engine)

	loginBody, err := json.Marshal(map[string]string{
		"username": "root",
		"password": "123456",
	})
	if err != nil {
		t.Fatalf("failed to marshal login body: %v", err)
	}
	loginReq := httptest.NewRequest(http.MethodPost, "/api/user/login", bytes.NewReader(loginBody))
	loginReq.Header.Set("Content-Type", "application/json")
	loginRecorder := httptest.NewRecorder()
	engine.ServeHTTP(loginRecorder, loginReq)
	if loginRecorder.Code != http.StatusOK {
		t.Fatalf("unexpected login status code: %d", loginRecorder.Code)
	}
	loginResult := loginRecorder.Result()
	defer loginResult.Body.Close()

	return engine, loginResult.Cookies()
}

func fakeManualServerBinary(version string) (string, []byte) {
	if runtime.GOOS == "windows" {
		return "atsflare-server-test.cmd", []byte("@echo off\r\necho " + version + "\r\n")
	}
	return "atsflare-server-test.sh", []byte("#!/bin/sh\necho " + version + "\n")
}

func TestManualUploadRoute(t *testing.T) {
	originalVersion := common.Version
	common.Version = "v0.4.0"
	t.Cleanup(func() {
		common.Version = originalVersion
		service.SetServerBinaryUpgradeExecutorForTest(nil)
		service.SetServerUpgradeDispatchDelayForTest(500 * time.Millisecond)
	})

	engine, cookies := loginRootAndBuildEngine(t)
	fileName, content := fakeManualServerBinary("v0.5.0")

	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)
	part, err := writer.CreateFormFile("binary", fileName)
	if err != nil {
		t.Fatalf("failed to create form file: %v", err)
	}
	if _, err = part.Write(content); err != nil {
		t.Fatalf("failed to write upload content: %v", err)
	}
	if err = writer.Close(); err != nil {
		t.Fatalf("failed to close multipart writer: %v", err)
	}

	req := httptest.NewRequest(http.MethodPost, "/api/update/manual-upload", body)
	req.Header.Set("Content-Type", writer.FormDataContentType())
	for _, cookieValue := range cookies {
		req.AddCookie(cookieValue)
	}

	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)
	if recorder.Code != http.StatusOK {
		t.Fatalf("unexpected status code: %d", recorder.Code)
	}

	var resp apiResponse
	if err = json.Unmarshal(recorder.Body.Bytes(), &resp); err != nil {
		t.Fatalf("failed to decode response: %v", err)
	}
	if !resp.Success {
		t.Fatalf("expected success response, got message: %s", resp.Message)
	}

	var data map[string]any
	if err = json.Unmarshal(resp.Data, &data); err != nil {
		t.Fatalf("failed to decode response data: %v", err)
	}
	if data["detected_version"] != "v0.5.0" {
		t.Fatalf("unexpected detected_version: %#v", data["detected_version"])
	}
	if data["ready_to_upgrade"] != true {
		t.Fatalf("expected ready_to_upgrade to be true: %#v", data["ready_to_upgrade"])
	}
	if data["upload_token"] == "" {
		t.Fatal("expected upload_token to be returned")
	}
}

func TestManualUpgradeConfirmRoute(t *testing.T) {
	originalVersion := common.Version
	originalExecutor := service.ServerBinaryUpgradeExecutorForTest()
	originalDelay := service.ServerUpgradeDispatchDelayForTest()
	common.Version = "v0.4.0"
	called := make(chan string, 1)
	service.SetServerBinaryUpgradeExecutorForTest(func(execPath string, tempPath string) error {
		called <- tempPath
		return nil
	})
	service.SetServerUpgradeDispatchDelayForTest(0)
	t.Cleanup(func() {
		common.Version = originalVersion
		service.SetServerBinaryUpgradeExecutorForTest(originalExecutor)
		service.SetServerUpgradeDispatchDelayForTest(originalDelay)
	})

	engine, cookies := loginRootAndBuildEngine(t)
	fileName, content := fakeManualServerBinary("v0.5.0")

	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)
	part, err := writer.CreateFormFile("binary", fileName)
	if err != nil {
		t.Fatalf("failed to create form file: %v", err)
	}
	if _, err = part.Write(content); err != nil {
		t.Fatalf("failed to write upload content: %v", err)
	}
	if err = writer.Close(); err != nil {
		t.Fatalf("failed to close multipart writer: %v", err)
	}

	uploadReq := httptest.NewRequest(http.MethodPost, "/api/update/manual-upload", body)
	uploadReq.Header.Set("Content-Type", writer.FormDataContentType())
	for _, cookieValue := range cookies {
		uploadReq.AddCookie(cookieValue)
	}

	uploadRecorder := httptest.NewRecorder()
	engine.ServeHTTP(uploadRecorder, uploadReq)
	if uploadRecorder.Code != http.StatusOK {
		t.Fatalf("unexpected upload status code: %d", uploadRecorder.Code)
	}

	var uploadResp apiResponse
	if err = json.Unmarshal(uploadRecorder.Body.Bytes(), &uploadResp); err != nil {
		t.Fatalf("failed to decode upload response: %v", err)
	}
	if !uploadResp.Success {
		t.Fatalf("expected upload success, got message: %s", uploadResp.Message)
	}

	var uploadData map[string]any
	if err = json.Unmarshal(uploadResp.Data, &uploadData); err != nil {
		t.Fatalf("failed to decode upload response data: %v", err)
	}
	uploadToken, _ := uploadData["upload_token"].(string)
	if uploadToken == "" {
		t.Fatal("expected upload token in upload response")
	}

	confirmBody, err := json.Marshal(map[string]string{"upload_token": uploadToken})
	if err != nil {
		t.Fatalf("failed to marshal confirm body: %v", err)
	}
	confirmReq := httptest.NewRequest(http.MethodPost, "/api/update/manual-upgrade", bytes.NewReader(confirmBody))
	confirmReq.Header.Set("Content-Type", "application/json")
	for _, cookieValue := range cookies {
		confirmReq.AddCookie(cookieValue)
	}

	confirmRecorder := httptest.NewRecorder()
	engine.ServeHTTP(confirmRecorder, confirmReq)
	if confirmRecorder.Code != http.StatusOK {
		t.Fatalf("unexpected confirm status code: %d", confirmRecorder.Code)
	}

	var confirmResp apiResponse
	if err = json.Unmarshal(confirmRecorder.Body.Bytes(), &confirmResp); err != nil {
		t.Fatalf("failed to decode confirm response: %v", err)
	}
	if !confirmResp.Success {
		t.Fatalf("expected confirm success, got message: %s", confirmResp.Message)
	}

	select {
	case tempPath := <-called:
		if tempPath == "" {
			t.Fatal("expected manual upgrade executor to receive temp path")
		}
	case <-time.After(time.Second):
		t.Fatal("expected manual upgrade executor to be called")
	}
}