From 27d25c68cccf8323ad0a6171710c1d7224e50d1c Mon Sep 17 00:00:00 2001
From: Cian McGrath <mcgrathc@arista.com>
Date: Wed, 27 Nov 2024 15:38:50 +0000
Subject: [PATCH] cloudvision/cvlib: Add context cleanup function

Adds a context cleanup function to close open connections.
This will be called at the end of a script/template execution
by the script executor infrastructure

Change-Id: I385cd9baf6e33ff904330af62fae5a96538df207
---
 cloudvision/cvlib/context.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/cloudvision/cvlib/context.py b/cloudvision/cvlib/context.py
index 5686be4..1e1ab4a 100644
--- a/cloudvision/cvlib/context.py
+++ b/cloudvision/cvlib/context.py
@@ -126,6 +126,16 @@ def __init__(self, user: User,
         self.studioCustomData = StudioCustomData(self)
         self.workspace: Optional[Workspace] = None
 
+    def cleanup(self):
+        '''
+        Cleans up open channels and other connections
+        Called by the script execution infra once a script or template has been executed
+        '''
+        if self.__connector:
+            self.__connector.close()
+        if self.__serviceChann:
+            self.__serviceChann.close()
+
     def getDevice(self):
         '''
         Returns the device associated to the context.