Merge pull request #30569 from ruchamahabal/fix-leave-alloc-update

diff --git a/erpnext/hr/doctype/leave_allocation/leave_allocation.py b/erpnext/hr/doctype/leave_allocation/leave_allocation.py
index 98408af..27479a5 100755
--- a/erpnext/hr/doctype/leave_allocation/leave_allocation.py
+++ b/erpnext/hr/doctype/leave_allocation/leave_allocation.py
@@ -39,11 +39,15 @@
 	def validate(self):
 		self.validate_period()
 		self.validate_allocation_overlap()
-		self.validate_back_dated_allocation()
-		self.set_total_leaves_allocated()
-		self.validate_total_leaves_allocated()
 		self.validate_lwp()
 		set_employee_name(self)
+		self.set_total_leaves_allocated()
+		self.validate_leave_days_and_dates()
+
+	def validate_leave_days_and_dates(self):
+		# all validations that should run on save as well as on update after submit
+		self.validate_back_dated_allocation()
+		self.validate_total_leaves_allocated()
 		self.validate_leave_allocation_days()
 
 	def validate_leave_allocation_days(self):
@@ -56,14 +60,19 @@
 			leave_allocated = 0
 			if leave_period:
 				leave_allocated = get_leave_allocation_for_period(
-					self.employee, self.leave_type, leave_period[0].from_date, leave_period[0].to_date
+					self.employee,
+					self.leave_type,
+					leave_period[0].from_date,
+					leave_period[0].to_date,
+					exclude_allocation=self.name,
 				)
 			leave_allocated += flt(self.new_leaves_allocated)
 			if leave_allocated > max_leaves_allowed:
 				frappe.throw(
 					_(
-						"Total allocated leaves are more days than maximum allocation of {0} leave type for employee {1} in the period"
-					).format(self.leave_type, self.employee)
+						"Total allocated leaves are more than maximum allocation allowed for {0} leave type for employee {1} in the period"
+					).format(self.leave_type, self.employee),
+					OverAllocationError,
 				)
 
 	def on_submit(self):
@@ -84,6 +93,12 @@
 	def on_update_after_submit(self):
 		if self.has_value_changed("new_leaves_allocated"):
 			self.validate_against_leave_applications()
+
+			# recalculate total leaves allocated
+			self.total_leaves_allocated = flt(self.unused_leaves) + flt(self.new_leaves_allocated)
+			# run required validations again since total leaves are being updated
+			self.validate_leave_days_and_dates()
+
 			leaves_to_be_added = self.new_leaves_allocated - self.get_existing_leave_count()
 			args = {
 				"leaves": leaves_to_be_added,
@@ -92,6 +107,7 @@
 				"is_carry_forward": 0,
 			}
 			create_leave_ledger_entry(self, args, True)
+			self.db_update()
 
 	def get_existing_leave_count(self):
 		ledger_entries = frappe.get_all(
@@ -279,27 +295,27 @@
 	)
 
 
-def get_leave_allocation_for_period(employee, leave_type, from_date, to_date):
-	leave_allocated = 0
-	leave_allocations = frappe.db.sql(
-		"""
-		select employee, leave_type, from_date, to_date, total_leaves_allocated
-		from `tabLeave Allocation`
-		where employee=%(employee)s and leave_type=%(leave_type)s
-			and docstatus=1
-			and (from_date between %(from_date)s and %(to_date)s
-				or to_date between %(from_date)s and %(to_date)s
-				or (from_date < %(from_date)s and to_date > %(to_date)s))
-	""",
-		{"from_date": from_date, "to_date": to_date, "employee": employee, "leave_type": leave_type},
-		as_dict=1,
-	)
+def get_leave_allocation_for_period(
+	employee, leave_type, from_date, to_date, exclude_allocation=None
+):
+	from frappe.query_builder.functions import Sum
 
-	if leave_allocations:
-		for leave_alloc in leave_allocations:
-			leave_allocated += leave_alloc.total_leaves_allocated
-
-	return leave_allocated
+	Allocation = frappe.qb.DocType("Leave Allocation")
+	return (
+		frappe.qb.from_(Allocation)
+		.select(Sum(Allocation.total_leaves_allocated).as_("total_allocated_leaves"))
+		.where(
+			(Allocation.employee == employee)
+			& (Allocation.leave_type == leave_type)
+			& (Allocation.docstatus == 1)
+			& (Allocation.name != exclude_allocation)
+			& (
+				(Allocation.from_date.between(from_date, to_date))
+				| (Allocation.to_date.between(from_date, to_date))
+				| ((Allocation.from_date < from_date) & (Allocation.to_date > to_date))
+			)
+		)
+	).run()[0][0] or 0.0
 
 
 @frappe.whitelist()
diff --git a/erpnext/hr/doctype/leave_allocation/test_leave_allocation.py b/erpnext/hr/doctype/leave_allocation/test_leave_allocation.py
index a53d4a8..dde52d7 100644
--- a/erpnext/hr/doctype/leave_allocation/test_leave_allocation.py
+++ b/erpnext/hr/doctype/leave_allocation/test_leave_allocation.py
@@ -1,24 +1,26 @@
 import unittest
 
 import frappe
+from frappe.tests.utils import FrappeTestCase
 from frappe.utils import add_days, add_months, getdate, nowdate
 
 import erpnext
 from erpnext.hr.doctype.employee.test_employee import make_employee
+from erpnext.hr.doctype.leave_allocation.leave_allocation import (
+	BackDatedAllocationError,
+	OverAllocationError,
+)
 from erpnext.hr.doctype.leave_ledger_entry.leave_ledger_entry import process_expired_allocation
 from erpnext.hr.doctype.leave_type.test_leave_type import create_leave_type
 
 
-class TestLeaveAllocation(unittest.TestCase):
-	@classmethod
-	def setUpClass(cls):
-		frappe.db.sql("delete from `tabLeave Period`")
+class TestLeaveAllocation(FrappeTestCase):
+	def setUp(self):
+		frappe.db.delete("Leave Period")
+		frappe.db.delete("Leave Allocation")
 
-		emp_id = make_employee("test_emp_leave_allocation@salary.com")
-		cls.employee = frappe.get_doc("Employee", emp_id)
-
-	def tearDown(self):
-		frappe.db.rollback()
+		emp_id = make_employee("test_emp_leave_allocation@salary.com", company="_Test Company")
+		self.employee = frappe.get_doc("Employee", emp_id)
 
 	def test_overlapping_allocation(self):
 		leaves = [
@@ -65,7 +67,7 @@
 		# invalid period
 		self.assertRaises(frappe.ValidationError, doc.save)
 
-	def test_allocated_leave_days_over_period(self):
+	def test_validation_for_over_allocation(self):
 		doc = frappe.get_doc(
 			{
 				"doctype": "Leave Allocation",
@@ -80,7 +82,135 @@
 		)
 
 		# allocated leave more than period
-		self.assertRaises(frappe.ValidationError, doc.save)
+		self.assertRaises(OverAllocationError, doc.save)
+
+	def test_validation_for_over_allocation_post_submission(self):
+		allocation = frappe.get_doc(
+			{
+				"doctype": "Leave Allocation",
+				"__islocal": 1,
+				"employee": self.employee.name,
+				"employee_name": self.employee.employee_name,
+				"leave_type": "_Test Leave Type",
+				"from_date": getdate("2015-09-1"),
+				"to_date": getdate("2015-09-30"),
+				"new_leaves_allocated": 15,
+			}
+		).submit()
+		allocation.reload()
+		# allocated leaves more than period after submission
+		allocation.new_leaves_allocated = 35
+		self.assertRaises(OverAllocationError, allocation.save)
+
+	def test_validation_for_over_allocation_based_on_leave_setup(self):
+		frappe.delete_doc_if_exists("Leave Period", "Test Allocation Period")
+		leave_period = frappe.get_doc(
+			dict(
+				name="Test Allocation Period",
+				doctype="Leave Period",
+				from_date=add_months(nowdate(), -6),
+				to_date=add_months(nowdate(), 6),
+				company="_Test Company",
+				is_active=1,
+			)
+		).insert()
+
+		leave_type = create_leave_type(leave_type_name="_Test Allocation Validation", is_carry_forward=1)
+		leave_type.max_leaves_allowed = 25
+		leave_type.save()
+
+		# 15 leaves allocated in this period
+		allocation = create_leave_allocation(
+			leave_type=leave_type.name,
+			employee=self.employee.name,
+			employee_name=self.employee.employee_name,
+			from_date=leave_period.from_date,
+			to_date=nowdate(),
+		)
+		allocation.submit()
+
+		# trying to allocate additional 15 leaves
+		allocation = create_leave_allocation(
+			leave_type=leave_type.name,
+			employee=self.employee.name,
+			employee_name=self.employee.employee_name,
+			from_date=add_days(nowdate(), 1),
+			to_date=leave_period.to_date,
+		)
+		self.assertRaises(OverAllocationError, allocation.save)
+
+	def test_validation_for_over_allocation_based_on_leave_setup_post_submission(self):
+		frappe.delete_doc_if_exists("Leave Period", "Test Allocation Period")
+		leave_period = frappe.get_doc(
+			dict(
+				name="Test Allocation Period",
+				doctype="Leave Period",
+				from_date=add_months(nowdate(), -6),
+				to_date=add_months(nowdate(), 6),
+				company="_Test Company",
+				is_active=1,
+			)
+		).insert()
+
+		leave_type = create_leave_type(leave_type_name="_Test Allocation Validation", is_carry_forward=1)
+		leave_type.max_leaves_allowed = 30
+		leave_type.save()
+
+		# 15 leaves allocated
+		allocation = create_leave_allocation(
+			leave_type=leave_type.name,
+			employee=self.employee.name,
+			employee_name=self.employee.employee_name,
+			from_date=leave_period.from_date,
+			to_date=nowdate(),
+		)
+		allocation.submit()
+		allocation.reload()
+
+		# allocate additional 15 leaves
+		allocation = create_leave_allocation(
+			leave_type=leave_type.name,
+			employee=self.employee.name,
+			employee_name=self.employee.employee_name,
+			from_date=add_days(nowdate(), 1),
+			to_date=leave_period.to_date,
+		)
+		allocation.submit()
+		allocation.reload()
+
+		# trying to allocate 25 leaves in 2nd alloc within leave period
+		# total leaves = 40 which is more than `max_leaves_allowed` setting i.e. 30
+		allocation.new_leaves_allocated = 25
+		self.assertRaises(OverAllocationError, allocation.save)
+
+	def test_validate_back_dated_allocation_update(self):
+		leave_type = create_leave_type(leave_type_name="_Test_CF_leave", is_carry_forward=1)
+		leave_type.save()
+
+		# initial leave allocation = 15
+		leave_allocation = create_leave_allocation(
+			employee=self.employee.name,
+			employee_name=self.employee.employee_name,
+			leave_type="_Test_CF_leave",
+			from_date=add_months(nowdate(), -12),
+			to_date=add_months(nowdate(), -1),
+			carry_forward=0,
+		)
+		leave_allocation.submit()
+
+		# new_leaves = 15, carry_forwarded = 10
+		leave_allocation_1 = create_leave_allocation(
+			employee=self.employee.name,
+			employee_name=self.employee.employee_name,
+			leave_type="_Test_CF_leave",
+			carry_forward=1,
+		)
+		leave_allocation_1.submit()
+
+		# try updating initial leave allocation
+		leave_allocation.reload()
+		leave_allocation.new_leaves_allocated = 20
+		self.assertRaises(BackDatedAllocationError, leave_allocation.save)
 
 	def test_carry_forward_calculation(self):
 		leave_type = create_leave_type(leave_type_name="_Test_CF_leave", is_carry_forward=1)
@@ -108,8 +238,10 @@
 			carry_forward=1,
 		)
 		leave_allocation_1.submit()
+		leave_allocation_1.reload()
 
 		self.assertEqual(leave_allocation_1.unused_leaves, 10)
+		self.assertEqual(leave_allocation_1.total_leaves_allocated, 25)
 
 		leave_allocation_1.cancel()
 
@@ -197,9 +329,12 @@
 			employee=self.employee.name, employee_name=self.employee.employee_name
 		)
 		leave_allocation.submit()
+		leave_allocation.reload()
 		self.assertTrue(leave_allocation.total_leaves_allocated, 15)
+
 		leave_allocation.new_leaves_allocated = 40
 		leave_allocation.submit()
+		leave_allocation.reload()
 		self.assertTrue(leave_allocation.total_leaves_allocated, 40)
 
 	def test_leave_subtraction_after_submit(self):
@@ -207,9 +342,12 @@
 			employee=self.employee.name, employee_name=self.employee.employee_name
 		)
 		leave_allocation.submit()
+		leave_allocation.reload()
 		self.assertTrue(leave_allocation.total_leaves_allocated, 15)
+
 		leave_allocation.new_leaves_allocated = 10
 		leave_allocation.submit()
+		leave_allocation.reload()
 		self.assertTrue(leave_allocation.total_leaves_allocated, 10)
 
 	def test_validation_against_leave_application_after_submit(self):