Added UTs for BountyBudgetWithdraw

stakwork · Mar 11, 2024 · 4063a97 · 4063a97
1 parent aeb39f9
commit 4063a97
Showing 1 changed file with 104 additions and 0 deletions.
diff --git a/handlers/bounty_test.go b/handlers/bounty_test.go
@@ -1567,6 +1567,110 @@ func TestBountyBudgetWithdraw(t *testing.T) {
 		assert.Equal(t, "Payment error", response["error"].(string))
 		mockHttpClient.AssertCalled(t, "Do", mock.AnythingOfType("*http.Request"))
 	})
+
+	t.Run("Should test that the BountyBudgetWithdraw handler gets locked by go mutex when it is called i.e. the handler has to be fully executed before it processes another request.", func(t *testing.T) {
+		ctx := context.Background()
+		authorizedCtx := context.WithValue(ctx, auth.ContextKey, "valid-key")
+		mockDb := dbMocks.NewDatabase(t)
+		mockHttpClient := mocks.NewHttpClient(t)
+		bHandler := NewBountyHandler(mockHttpClient, mockDb)
+		paymentAmount := uint(1500)
+		mockDb.On("UserHasAccess", "valid-key", mock.AnythingOfType("string"), db.WithdrawBudget).Return(true)
+		mockDb.On("GetOrganizationBudget", "org-1").Return(db.BountyBudget{TotalBudget: 3000}, nil)
+		mockDb.On("WithdrawBudget", "valid-key", "org-1", paymentAmount).Return(nil)
+		mockHttpClient.On("Do", mock.AnythingOfType("*http.Request")).Return(&http.Response{
+			StatusCode: 200,
+			Body:       io.NopCloser(bytes.NewBufferString(`{"success": true}`)),
+		}, nil)
+
+		invoice := "lnbc15u1p3xnhl2pp5jptserfk3zk4qy42tlucycrfwxhydvlemu9pqr93tuzlv9cc7g3sdqsvfhkcap3xyhx7un8cqzpgxqzjcsp5f8c52y2stc300gl6s4xswtjpc37hrnnr3c9wvtgjfuvqmpm35evq9qyyssqy4lgd8tj637qcjp05rdpxxykjenthxftej7a2zzmwrmrl70fyj9hvj0rewhzj7jfyuwkwcg9g2jpwtk3wkjtwnkdks84hsnu8xps5vsq4gj5hs"
+
+		withdrawRequest := db.WithdrawBudgetRequest{
+			PaymentRequest: invoice,
+			OrgUuid:        "org-1",
+		}
+
+		var wg sync.WaitGroup
+		var mu sync.Mutex
+		callTimes := make([]time.Time, 0, 2)
+
+		for i := 0; i < 2; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				requestBody, _ := json.Marshal(withdrawRequest)
+				req, err := http.NewRequestWithContext(authorizedCtx, http.MethodPost, "/budget/withdraw", bytes.NewReader(requestBody))
+				assert.NoError(t, err)
+				rr := httptest.NewRecorder()
+
+				mu.Lock()
+				bHandler.BountyBudgetWithdraw(rr, req)
+				time.Sleep(10 * time.Millisecond)
+				mu.Unlock()
+				callTimes = append(callTimes, time.Now())
+			}()
+		}
+
+		wg.Wait()
+
+		assert.Equal(t, 2, len(callTimes))
+		assert.True(t, callTimes[1].After(callTimes[0]), "The second request should be processed after the first one finishes")
+	})
+
+	t.Run("Should test that an Organization's Budget Total Amount is accurate after three (3) successful 'Budget Withdrawal Requests'", func(t *testing.T) {
+		ctxs := context.WithValue(context.Background(), auth.ContextKey, "valid-key")
+		mockDb := dbMocks.NewDatabase(t)
+		mockHttpClient := mocks.NewHttpClient(t)
+		bHandler := NewBountyHandler(mockHttpClient, mockDb)
+
+		paymentAmount := uint(1500)
+		initialBudget := uint(5000)
+		invoices := []string{
+			"lnbc15u1pj773m8pp5pj0neu674ka9g75tzc9em2crjnzdn2tfjplwua5dhzr2dcg5mrmsdqhgf6kgem9wssyjmnkda5kxegcqzpgxqyz5vqrzjqwnw5tv745sjpvft6e3f9w62xqk826vrm3zaev4nvj6xr3n065aukqqqqyqqqqsqqyqqqqqqqqqqqqqqqqsp5cg7rnce0u7sgqs5ayws0nuapcndps2uecgmflcrdvhs86cvr3spq9qrssql866kyj3cmcjts3mv927p92zaf0yhpl4uxnj7c429308gr4dwaxkh5rnhksuvrxd96dja60gl8p44rnuavszvy2p8h7a9ak7cllfrjcq2t04yg",
+			"lnbc15u1pj773a7pp5d934ereytmkh83f88fd5fn0388sfj54xdru2tcgmzlywmk3x0jvqdqhgf6kgem9wssyjmnkda5kxegcqzpgxqyz5vqrzjqwnw5tv745sjpvft6e3f9w62xqk826vrm3zaev4nvj6xr3n065aukqqqqyqqqqsqqyqqqqqqqqqqqqqqqqsp5dgn9x240n5pn2r8hazc9udxpg8tsrl8kwud2xjphx3f7natn7dws9qrssq64rhcg9syd90jxc8vz9l6xlfku7vyq8p20s9fjydggqh48ptkz9rwnm4f8qmsntjqqa3qv5qh7mm9gm2efxeqph5p93zgn2cs708pgcpy9dlue",
+			"lnbc15u1pj773ldpp5cgps7hk2cnwnj85a50zzgcap32srmd2zh20spek27mfgwddv9vksdqhgf6kgem9wssyjmnkda5kxegcqzpgxqyz5vqrzjqwnw5tv745sjpvft6e3f9w62xqk826vrm3zaev4nvj6xr3n065aukqqqqyqqqqsqqyqqqqqqqqqqqqqqqqsp5z8u0laltkm4dchcrd66pyx6x7a0lawqludzz9lr0pplkdtj5k5ds9qrssqt4vj5q9k3nffjqdsgw3dd83reu93k6skyydw2tst3vmuxvsyu5f5cftt2ucmmj0r2535tzsyx6rwuav4fgynmsdx6jrlvd7zw54tkccprxkzl2",
+		}
+
+		for i := 0; i < 3; i++ {
+			expectedFinalBudget := initialBudget - (paymentAmount * uint(i))
+
+			mockDb.ExpectedCalls = nil
+			mockDb.Calls = nil
+			mockHttpClient.ExpectedCalls = nil
+			mockHttpClient.Calls = nil
+
+			mockDb.On("UserHasAccess", "valid-key", "org-1", db.WithdrawBudget).Return(true)
+			//mockDb.On("GetOrganizationBudget", "org-1").Return(db.BountyBudget{TotalBudget: initialBudget - uint(i)*paymentAmount}, nil)
+			mockDb.On("GetOrganizationBudget", "org-1").Return(db.BountyBudget{
+				TotalBudget: expectedFinalBudget,
+			}, nil)
+			mockDb.On("WithdrawBudget", "valid-key", "org-1", paymentAmount).Return(nil)
+			mockHttpClient.On("Do", mock.AnythingOfType("*http.Request")).Return(&http.Response{
+				StatusCode: 200,
+				Body:       io.NopCloser(bytes.NewBufferString(`{"success": true}`)),
+			}, nil)
+
+			withdrawRequest := db.WithdrawBudgetRequest{
+				PaymentRequest: invoices[i],
+				OrgUuid:        "org-1",
+			}
+			requestBody, _ := json.Marshal(withdrawRequest)
+			req, _ := http.NewRequestWithContext(ctxs, http.MethodPost, "/budget/withdraw", bytes.NewReader(requestBody))
+
+			rr := httptest.NewRecorder()
+
+			bHandler.BountyBudgetWithdraw(rr, req)
+			assert.Equal(t, http.StatusOK, rr.Code)
+			var response db.InvoicePaySuccess
+			err := json.Unmarshal(rr.Body.Bytes(), &response)
+			assert.NoError(t, err)
+			assert.True(t, response.Success, "Expected invoice payment to succeed")
+			finalBudget := mockDb.GetOrganizationBudget("org-1")
+			assert.Equal(t, expectedFinalBudget, finalBudget.TotalBudget, "The organization's final budget should reflect the deductions from the successful withdrawals")
+
+		}
+	})
+
 }
 
 func TestPollInvoice(t *testing.T) {