From d16ef3fa072256fda619e8797576a9bb28c89554 Mon Sep 17 00:00:00 2001
From: ImMin5 <mino@megazone.com>
Date: Fri, 12 Jan 2024 17:11:22 +0900
Subject: [PATCH] fix: fix budget update error because of datasource_id
 parameter

---
 .../manager/budget_usage_manager.py           | 25 +++----------------
 .../cost_analysis/service/budget_service.py   | 15 +++--------
 2 files changed, 7 insertions(+), 33 deletions(-)

diff --git a/src/spaceone/cost_analysis/manager/budget_usage_manager.py b/src/spaceone/cost_analysis/manager/budget_usage_manager.py
index b6b9ff25..d6d29ae8 100644
--- a/src/spaceone/cost_analysis/manager/budget_usage_manager.py
+++ b/src/spaceone/cost_analysis/manager/budget_usage_manager.py
@@ -95,9 +95,7 @@ def update_cost_usage(
 
         self._update_monthly_budget_usage(budget_vo, cost_mgr)
 
-    def update_budget_usage(
-        self, domain_id: str, data_source_id: str
-    ):
+    def update_budget_usage(self, domain_id: str, data_source_id: str):
         budget_vos = self.budget_mgr.filter_budgets(
             domain_id=domain_id,
             data_source_id=data_source_id,
@@ -325,7 +323,8 @@ def _update_monthly_budget_usage(self, budget_vo: Budget, cost_mgr: CostManager)
             else:
                 budget_usage_vo.update({"cost": 0})
 
-    def _make_cost_analyze_query(self, budget_vo: Budget):
+    @staticmethod
+    def _make_cost_analyze_query(budget_vo: Budget):
         query = {
             "granularity": "MONTHLY",
             "start": budget_vo.start,
@@ -342,24 +341,6 @@ def _make_cost_analyze_query(self, budget_vo: Budget):
             query["filter"].append(
                 {"k": "project_id", "v": budget_vo.project_id, "o": "eq"}
             )
-        # else:
-        #     identity_mgr: IdentityManager = self.locator.get_manager("IdentityManager")
-        #
-        #     project_query = {
-        #         "filter": [
-        #             {"k": "domain_id", "v": budget_vo.domain_id, "o": "eq"},
-        #             {"k": "workspace_id", "v": budget_vo.workspace_id, "o": "eq"},
-        #         ]
-        #     }
-        #     projects_info = identity_mgr.list_projects(
-        #         project_query, budget_vo.domain_id
-        #     )
-        #
-        #     project_ids = []
-        #     for project_info in projects_info.get("results", []):
-        #         project_ids.append(project_info["project_id"])
-        #
-        #     query["filter"].append({"k": "project_id", "v": project_ids, "o": "in"})
 
         if budget_vo.provider_filter and budget_vo.provider_filter.state == "ENABLED":
             query["filter"].append(
diff --git a/src/spaceone/cost_analysis/service/budget_service.py b/src/spaceone/cost_analysis/service/budget_service.py
index d5a2b0c6..702a52a2 100644
--- a/src/spaceone/cost_analysis/service/budget_service.py
+++ b/src/spaceone/cost_analysis/service/budget_service.py
@@ -78,7 +78,6 @@ def create(self, params):
         provider_filter_state = provider_filter.get("state", "DISABLED")
         notifications = params.get("notifications", [])
         resource_group = params["resource_group"]
-        # self._check_target(project_id, project_group_id, domain_id)
         self._check_time_period(start, end)
 
         identity_mgr: IdentityManager = self.locator.get_manager("IdentityManager")
@@ -149,9 +148,7 @@ def create(self, params):
             "BudgetUsageManager"
         )
         budget_usage_mgr.create_budget_usages(budget_vo)
-        budget_usage_mgr.update_cost_usage(
-            budget_vo=budget_vo, data_source_workspace_id=data_source_vo.workspace_id
-        )
+        budget_usage_mgr.update_cost_usage(budget_vo)
         budget_usage_mgr.notify_budget_usage(budget_vo)
 
         return budget_vo
@@ -194,7 +191,6 @@ def update(self, params):
         )
 
         # Check limit and Planned Limits
-
         budget_vo = self.budget_mgr.update_budget_by_vo(params, budget_vo)
 
         if "name" in params:
@@ -206,14 +202,11 @@ def update(self, params):
                 budget_usage_mgr.update_budget_usage_by_vo(
                     {"name": params["name"]}, budget_usage_vo
                 )
+        # Check DataSource exists
         data_source_mgr = self.locator.get_manager("DataSourceManager")
-        data_source_vo = data_source_mgr.get_data_source(
-            budget_vo.data_source_id, domain_id
-        )
+        data_source_mgr.get_data_source(budget_vo.data_source_id, domain_id)
 
-        budget_usage_mgr.update_cost_usage(
-            budget_vo=budget_vo, data_source_workspace_id=data_source_vo.workspace_id
-        )
+        budget_usage_mgr.update_cost_usage(budget_vo)
         budget_usage_mgr.notify_budget_usage(budget_vo)
 
         return budget_vo