Coverage for backend/core/auxiliary/models/ObjectTypeCounter.py: 96%

1# core/auxiliary/models/ObjectTypeCounter.py 

2from django.db import models 

3from django.db import transaction 

4from django.db.models import F 

5 

6from core.auxiliary.models.Flowsheet import Flowsheet 

7from core.managers import AccessControlManager 

8 

9class ObjectTypeCounter(models.Model): 

10 """ 

11 Keeps a per-type counter for objects inside a flowsheet. 

12 Example: Heater1, Heater2, ST1, ST2, etc. 

13 """ 

14 objects = AccessControlManager() 

15 flowsheet = models.ForeignKey( 

16 Flowsheet, on_delete=models.CASCADE, related_name="type_counters" 

17 ) 

18 object_type = models.CharField(max_length=64) # e.g., "heater", "compressor", "stream" 

19 next_index = models.IntegerField(default=1) 

20 

21 class Meta: 

22 unique_together = ("flowsheet", "object_type") # ensures one counter per (flowsheet, type) 

23 

24 def __str__(self): 

25 return f"{self.flowsheet.name} - {self.object_type}: {self.next_index}" 

26 

27 @classmethod 

28 def next_for(cls, flowsheet: Flowsheet, object_type: str) -> int: 

29 """ 

30 Atomically reserve and return the next index for (flowsheet, object_type). 

31 Safe even when multiple users create objects at the same time. 

32 """ 

33 with transaction.atomic(): 

34 counter, _ = cls.objects.select_for_update().get_or_create( 

35 flowsheet=flowsheet, 

36 object_type=object_type, 

37 defaults={"next_index": 1}, 

38 ) 

39 current = counter.next_index 

40 counter.next_index = F("next_index") + 1 

41 counter.save(update_fields=["next_index"]) 

42 counter.refresh_from_db(fields=["next_index"]) 

43 return current