Coverage for backend/core/viewset.py: 92%

1from rest_framework import viewsets 

2 

3 

4class ContextMixin: 

5 """ 

6 A mixin that provides the flowsheet context to all views to enforce access control. 

7 """ 

8 def initial(self, request, *args, **kwargs): 

9 if hasattr(self, 'action'): 9 ↛ 14line 9 didn't jump to line 14 because the condition on line 9 was always true

10 view_method = getattr(self, self.action, None) 

11 if hasattr(view_method, 'ignore_access_control') and view_method.ignore_access_control: 

12 return super().initial(request, *args, **kwargs) 

13 

14 if hasattr(request, "query_params"): 14 ↛ 17line 14 didn't jump to line 17 because the condition on line 14 was always true

15 flowsheet_id = request.query_params.get("flowsheet") 

16 else: 

17 flowsheet_id = request.GET.get("flowsheet") 

18 

19 from core.validation import sanitize_flowsheet_id, flowsheet_context 

20 

21 sanitize_flowsheet_id(flowsheet_id) 

22 self.flowsheet_context = flowsheet_context(flowsheet_id, request.user) 

23 # start the context 

24 self.flowsheet_token = self.flowsheet_context.__enter__() 

25 

26 return super().initial(request, *args, **kwargs) 

27 

28 def finalize_response(self, request, response, *args, **kwargs): 

29 # Clean up context when the response is finalized 

30 if hasattr(self, 'flowsheet_context'): 

31 self.flowsheet_context.__exit__(None, None, None) 

32 

33 return super().finalize_response(request, response, *args, **kwargs) 

34 

35 def handle_exception(self, exc): 

36 # Clean up context in case of exception 

37 if hasattr(self, 'flowsheet_context'): 

38 self.flowsheet_context.__exit__(type(exc), exc, exc.__traceback__) 

39 

40 return super().handle_exception(exc) 

41 

42class ReadOnlyModelViewSet(ContextMixin, viewsets.ReadOnlyModelViewSet): 

43 """ 

44 Use this model viewset for read-only operations instead of the default one to enforce access control. 

45 """ 

46 pass 

47 

48class ModelViewSet(ContextMixin, viewsets.ModelViewSet): 

49 """ 

50 Use this model viewset instead of the default one to enforce access control. 

51 """ 

52 pass