From 3eda8308e1520ab544a550337a38c28dc948534f Mon Sep 17 00:00:00 2001
From: dtfiedler <dylan@ar.io>
Date: Fri, 18 Oct 2024 14:05:19 -0500
Subject: [PATCH] chore(test): add tests for getTokenCost

---
 spec/arns_spec.lua | 50 ++++++++++++++++++++++++++++++++++++++++++++++
 src/arns.lua       |  9 +++++----
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/spec/arns_spec.lua b/spec/arns_spec.lua
index e0b4611..6d720a8 100644
--- a/spec/arns_spec.lua
+++ b/spec/arns_spec.lua
@@ -424,6 +424,56 @@ describe("arns", function()
 		end)
 	end
 
+	describe("getTokenCost", function()
+		it("should return the correct token cost for a lease", function()
+			local baseFee = 500000000
+			local years = 2
+			local demandFactor = 1.0
+			local expectedCost = ((years * baseFee * 0.20) + baseFee) * demandFactor
+			local intendedAction = {
+				intent = "Buy-Record",
+				purchaseType = "lease",
+				years = 2,
+				name = "test-name",
+			}
+			assert.are.equal(expectedCost, arns.getTokenCost(intendedAction))
+		end)
+		it("should return the correct token cost for a permabuy", function()
+			local baseFee = 500000000
+			local demandFactor = 1.0
+			local expectedCost = ((baseFee * 0.2 * 20) + baseFee) * demandFactor
+			local intendedAction = {
+				intent = "Buy-Record",
+				purchaseType = "permabuy",
+				name = "test-name",
+			}
+			assert.are.equal(expectedCost, arns.getTokenCost(intendedAction))
+		end)
+		it("should return the correct token cost for an undername", function()
+			_G.NameRegistry.records["test-name"] = {
+				endTimestamp = constants.oneYearMs,
+				processId = testProcessId,
+				purchasePrice = 600000000,
+				startTimestamp = 0,
+				type = "lease",
+				undernameLimit = 10,
+			}
+			local baseFee = 500000000
+			local undernamePercentageFee = 0.001
+			local increaseQty = 5
+			local demandFactor = 1.0
+			local yearsRemaining = 0.5
+			local expectedCost = baseFee * increaseQty * undernamePercentageFee * yearsRemaining * demandFactor
+			local intendedAction = {
+				intent = "Increase-Undername-Limit",
+				quantity = 5,
+				name = "test-name",
+				currentTimestamp = constants.oneYearMs / 2,
+			}
+			assert.are.equal(expectedCost, arns.getTokenCost(intendedAction))
+		end)
+	end)
+
 	describe("pruneRecords", function()
 		it("should prune records", function()
 			local currentTimestamp = 1000000000
diff --git a/src/arns.lua b/src/arns.lua
index ad64add..76c68ae 100644
--- a/src/arns.lua
+++ b/src/arns.lua
@@ -249,8 +249,7 @@ function arns.calculateLeaseFee(baseFee, years, demandFactor)
 end
 
 function arns.calculateAnnualRenewalFee(baseFee, years)
-	local nameAnnualRegistrationFee = baseFee * constants.ANNUAL_PERCENTAGE_FEE
-	local totalAnnualRenewalCost = nameAnnualRegistrationFee * years
+	local totalAnnualRenewalCost = baseFee * constants.ANNUAL_PERCENTAGE_FEE * years
 	return math.floor(totalAnnualRenewalCost)
 end
 
@@ -369,9 +368,11 @@ function arns.getTokenCost(intendedAction)
 		local purchaseType = intendedAction.purchaseType
 		local years = intendedAction.years
 		local name = intendedAction.name
-		assert(purchaseType == "lease" or purchaseType == "permabuy", "PurchaseType is invalid.")
-		assert(years >= 1 and years <= 5, "Years is invalid. Must be an integer between 1 and 5")
 		assert(type(name) == "string", "Name is required and must be a string.")
+		assert(purchaseType == "lease" or purchaseType == "permabuy", "PurchaseType is invalid.")
+		if purchaseType == "lease" then
+			assert(years >= 1 and years <= 5, "Years is invalid. Must be an integer between 1 and 5")
+		end
 		local baseFee = demand.getFees()[#name]
 		tokenCost = arns.calculateRegistrationFee(purchaseType, baseFee, years, demand.getDemandFactor())
 	elseif intendedAction.intent == "Extend-Lease" then