From c21b87db6bd330906116d31f4e308b2eced591e7 Mon Sep 17 00:00:00 2001
From: avi robusta <avi@robusta.dev>
Date: Thu, 18 May 2023 11:00:53 +0300
Subject: PR fixes

---
 robusta_krr/core/runner.py        | 15 ++++++++-------
 robusta_krr/utils/progress_bar.py | 12 ++++--------
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/robusta_krr/core/runner.py b/robusta_krr/core/runner.py
index 098f86e..b0779c2 100644
--- a/robusta_krr/core/runner.py
+++ b/robusta_krr/core/runner.py
@@ -91,7 +91,7 @@ class Runner(Configurable):
             for resource, recommendation in result.items()
         }
 
-    async def _calculate_object_recommendations(self, object: K8sObjectData, progress_bar: ProgressBar) -> RunResult:
+    async def _calculate_object_recommendations(self, object: K8sObjectData) -> RunResult:
         prometheus_loader = self._get_prometheus_loader(object.cluster)
 
         if prometheus_loader is None:
@@ -108,15 +108,15 @@ class Runner(Configurable):
             ]
         )
         data = dict(zip(ResourceType, data_tuple))
-        progress_bar.progress()
+        self.__progressbar.progress()
         # NOTE: We run this in a threadpool as the strategy calculation might be CPU intensive
         # But keep in mind that numpy calcluations will not block the GIL
         result = await asyncio.to_thread(self._strategy.run, data, object)
         return self._format_result(result)
 
-    async def _gather_objects_recommendations(self, objects: list[K8sObjectData], progress_bar: ProgressBar) -> list[ResourceAllocations]:
+    async def _gather_objects_recommendations(self, objects: list[K8sObjectData]) -> list[ResourceAllocations]:
         recommendations: list[RunResult] = await asyncio.gather(
-            *[self._calculate_object_recommendations(object, progress_bar) for object in objects]
+            *[self._calculate_object_recommendations(object) for object in objects]
         )
 
         return [
@@ -139,9 +139,10 @@ class Runner(Configurable):
                 self.warning("Note that you are using the '*' namespace filter, which by default excludes kube-system.")
             return Result(scans=[])
 
-        with ProgressBar(self.config, total=len(objects), title="Calculating strategy") as progress_bar:
-            resource_recommendations = await self._gather_objects_recommendations(objects, progress_bar)
-
+        self.__progressbar = ProgressBar(self.config, total=len(objects), title="Calculating Recommendation")
+        resource_recommendations = await self._gather_objects_recommendations(objects)
+        self.__progressbar.close_bar()
+        
         return Result(
             scans=[
                 ResourceScan.calculate(obj, recommended) for obj, recommended in zip(objects, resource_recommendations)
diff --git a/robusta_krr/utils/progress_bar.py b/robusta_krr/utils/progress_bar.py
index cd5038c..5bd71d5 100644
--- a/robusta_krr/utils/progress_bar.py
+++ b/robusta_krr/utils/progress_bar.py
@@ -1,7 +1,7 @@
 from robusta_krr.utils.configurable import Configurable
 from alive_progress import alive_bar
 from robusta_krr.core.models.config import Config
-
+import sys
 
 class ProgressBar(Configurable):
     def __init__(self, config: Config, **kwargs) -> None:
@@ -9,16 +9,12 @@ class ProgressBar(Configurable):
         self.show_bar = self.echo_active
         if self.show_bar:
             self.alive_bar = alive_bar(**kwargs)
-
-    def __enter__(self):
-        if self.show_bar:
             self.bar = self.alive_bar.__enter__()
-        return self
 
     def progress(self):
         if self.show_bar:
             self.bar()
-
-    def __exit__(self, *args):
+     
+    def close_bar(self):
         if self.show_bar:
-            self.alive_bar.__exit__(*args)
\ No newline at end of file
+            self.alive_bar.__exit__(*sys.exc_info())
\ No newline at end of file
-- 
cgit v1.2.3