smart-search-back/tests/transaction_rollback_test.go

package tests

import (
	"sync"
	"sync/atomic"

	authpb "git.techease.ru/Smart-search/smart-search-back/pkg/pb/auth"
	requestpb "git.techease.ru/Smart-search/smart-search-back/pkg/pb/request"
)

func (s *IntegrationSuite) TestTransaction_CreateTZ_InsufficientBalance_Rollback() {
	email, password, userID := s.createUniqueTestUser("insufficient_tz", 0.001)

	loginResp, err := s.authClient.Login(s.ctx, &authpb.LoginRequest{
		Email:     email,
		Password:  password,
		Ip:        "127.0.0.1",
		UserAgent: "test-agent",
	})
	s.Require().NoError(err)

	validateResp, err := s.authClient.Validate(s.ctx, &authpb.ValidateRequest{
		AccessToken: loginResp.AccessToken,
	})
	s.Require().NoError(err)

	initialBalance := s.getUserBalance(userID)
	s.T().Logf("Initial balance: %.4f", initialBalance)

	_, err = s.requestClient.CreateTZ(s.ctx, &requestpb.CreateTZRequest{
		UserId:     validateResp.UserId,
		RequestTxt: "Тест с недостаточным балансом",
	})

	if err != nil {
		s.T().Logf("CreateTZ failed as expected: %v", err)

		finalBalance := s.getUserBalance(userID)
		s.T().Logf("Final balance after failed CreateTZ: %.4f", finalBalance)

		s.GreaterOrEqual(finalBalance, 0.0,
			"Баланс не должен быть отрицательным после rollback")
	}
}

func (s *IntegrationSuite) TestTransaction_ApproveTZ_InsufficientBalance_NoSuppliers() {
	email, password, userID := s.createUniqueTestUser("approve_insufficient", 1000.0)

	loginResp, err := s.authClient.Login(s.ctx, &authpb.LoginRequest{
		Email:     email,
		Password:  password,
		Ip:        "127.0.0.1",
		UserAgent: "test-agent",
	})
	s.Require().NoError(err)

	validateResp, err := s.authClient.Validate(s.ctx, &authpb.ValidateRequest{
		AccessToken: loginResp.AccessToken,
	})
	s.Require().NoError(err)

	createResp, err := s.requestClient.CreateTZ(s.ctx, &requestpb.CreateTZRequest{
		UserId:     validateResp.UserId,
		RequestTxt: "Тест approve с недостаточным балансом",
	})
	s.Require().NoError(err)
	requestID := createResp.RequestId

	_, err = s.pool.Exec(s.ctx, "UPDATE users SET balance = 0.001 WHERE id = $1", userID)
	s.Require().NoError(err)

	suppliersBeforeApprove := s.getRequestSuppliersCount(requestID)
	s.T().Logf("Suppliers before ApproveTZ: %d", suppliersBeforeApprove)

	_, err = s.requestClient.ApproveTZ(s.ctx, &requestpb.ApproveTZRequest{
		RequestId: requestID,
		FinalTz:   "Утвержденное ТЗ",
		UserId:    validateResp.UserId,
	})

	if err != nil {
		s.T().Logf("ApproveTZ failed as expected: %v", err)

		suppliersAfterApprove := s.getRequestSuppliersCount(requestID)
		s.T().Logf("Suppliers after failed ApproveTZ: %d", suppliersAfterApprove)

		finalBalance := s.getUserBalance(userID)
		s.GreaterOrEqual(finalBalance, 0.0,
			"Баланс не должен быть отрицательным")
	}
}

func (s *IntegrationSuite) TestTransaction_ConcurrentCreateTZ_BalanceAtomicity() {
	email, password, userID := s.createUniqueTestUser("concurrent_tz", 100.0)

	loginResp, err := s.authClient.Login(s.ctx, &authpb.LoginRequest{
		Email:     email,
		Password:  password,
		Ip:        "127.0.0.1",
		UserAgent: "test-agent",
	})
	s.Require().NoError(err)

	validateResp, err := s.authClient.Validate(s.ctx, &authpb.ValidateRequest{
		AccessToken: loginResp.AccessToken,
	})
	s.Require().NoError(err)

	var wg sync.WaitGroup
	var successCount int32
	var errorCount int32
	goroutines := 10

	startBarrier := make(chan struct{})

	for i := 0; i < goroutines; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()

			<-startBarrier

			_, err := s.requestClient.CreateTZ(s.ctx, &requestpb.CreateTZRequest{
				UserId:     validateResp.UserId,
				RequestTxt: "Параллельный тест CreateTZ",
			})

			if err == nil {
				atomic.AddInt32(&successCount, 1)
			} else {
				atomic.AddInt32(&errorCount, 1)
			}
		}(i)
	}

	close(startBarrier)
	wg.Wait()

	s.T().Logf("Concurrent CreateTZ - Success: %d, Errors: %d", successCount, errorCount)

	finalBalance := s.getUserBalance(userID)
	s.T().Logf("Final balance: %.4f", finalBalance)

	s.GreaterOrEqual(finalBalance, 0.0,
		"Баланс не должен быть отрицательным после параллельных операций")
}

func (s *IntegrationSuite) TestTransaction_TokenUsage_BalanceConsistency() {
	email, password, userID := s.createUniqueTestUser("token_consistency", 1000.0)

	initialBalance := s.getUserBalance(userID)
	s.T().Logf("Initial balance: %.4f", initialBalance)

	loginResp, err := s.authClient.Login(s.ctx, &authpb.LoginRequest{
		Email:     email,
		Password:  password,
		Ip:        "127.0.0.1",
		UserAgent: "test-agent",
	})
	s.Require().NoError(err)

	validateResp, err := s.authClient.Validate(s.ctx, &authpb.ValidateRequest{
		AccessToken: loginResp.AccessToken,
	})
	s.Require().NoError(err)

	createResp, err := s.requestClient.CreateTZ(s.ctx, &requestpb.CreateTZRequest{
		UserId:     validateResp.UserId,
		RequestTxt: "Тест consistency token_usage и balance",
	})
	s.Require().NoError(err)
	requestID := createResp.RequestId

	tokenUsageCount := s.getTokenUsageCount(requestID)
	s.T().Logf("Token usage records for request: %d", tokenUsageCount)

	finalBalance := s.getUserBalance(userID)
	balanceDelta := initialBalance - finalBalance
	s.T().Logf("Balance delta: %.4f", balanceDelta)

	if tokenUsageCount > 0 {
		s.Greater(balanceDelta, 0.0,
			"Баланс должен уменьшиться при наличии token_usage записей")
	}

	var totalTokenCost float64
	err = s.pool.QueryRow(s.ctx,
		"SELECT COALESCE(SUM(token_cost), 0) FROM request_token_usage WHERE request_id = $1::uuid",
		requestID,
	).Scan(&totalTokenCost)
	s.NoError(err)

	s.T().Logf("Total token cost from DB: %.4f, Balance delta: %.4f", totalTokenCost, balanceDelta)
}