move update to Variable definition

statisticsnorway · Dec 13, 2024 · 524b36b · 524b36b
1 parent d8dcd32
commit 524b36b
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 72 deletions.
diff --git a/src/dapla_metadata/variable_definitions/vardef.py b/src/dapla_metadata/variable_definitions/vardef.py
@@ -15,9 +15,6 @@
 from dapla_metadata.variable_definitions.generated.vardef_client.models.draft import (
     Draft,
 )
-from dapla_metadata.variable_definitions.generated.vardef_client.models.update_draft import (
-    UpdateDraft,
-)
 from dapla_metadata.variable_definitions.variable_definition import VariableDefinition
 
 
@@ -94,44 +91,6 @@ def create_draft(cls, draft: Draft) -> VariableDefinition:
             ),
         )
 
-    @classmethod
-    @vardef_exception_handler
-    def update_draft(
-        cls,
-        variable_definition_id: str,
-        update_draft: UpdateDraft,
-    ) -> VariableDefinition:
-        """Update a Draft variable definition."""
-        return VariableDefinition.from_complete_response(
-            DraftVariableDefinitionsApi(
-                VardefClient.get_client(),
-            ).update_variable_definition_by_id(
-                active_group=config.get_active_group(),
-                variable_definition_id=variable_definition_id,
-                update_draft=update_draft,
-            ),
-        )
-
-    @classmethod
-    @vardef_exception_handler
-    def update_draft_by_short_name(
-        cls,
-        short_name: str,
-        update_draft: UpdateDraft,
-    ) -> VariableDefinition:
-        """Update a Draft variable definition."""
-        return VariableDefinition.from_complete_response(
-            DraftVariableDefinitionsApi(
-                VardefClient.get_client(),
-            ).update_variable_definition_by_id(
-                active_group=config.get_active_group(),
-                variable_definition_id=VariableDefinition.get_id_from_short_name(
-                    short_name=short_name,
-                ),
-                update_draft=update_draft,
-            ),
-        )
-
     @classmethod
     def migrate_from_vardok(cls, vardok_id: str) -> VariableDefinition:
         """Migrate a Variable Definition from Vardok to Vardef.

diff --git a/src/dapla_metadata/variable_definitions/variable_definition.py b/src/dapla_metadata/variable_definitions/variable_definition.py
@@ -1,18 +1,20 @@
-from tests.utils.constants import VARDEF_EXAMPLE_DATE
-
 from dapla_metadata.variable_definitions._client import VardefClient
+from dapla_metadata.variable_definitions.exceptions import vardef_exception_handler
+from dapla_metadata.variable_definitions.generated.vardef_client.api.draft_variable_definitions_api import (
+    DraftVariableDefinitionsApi,
+)
 from dapla_metadata.variable_definitions.generated.vardef_client.api.patches_api import (
     PatchesApi,
 )
 from dapla_metadata.variable_definitions.generated.vardef_client.api.validity_periods_api import (
     ValidityPeriodsApi,
 )
-from dapla_metadata.variable_definitions.generated.vardef_client.api.variable_definitions_api import (
-    VariableDefinitionsApi,
-)
 from dapla_metadata.variable_definitions.generated.vardef_client.models.complete_response import (
     CompleteResponse,
 )
+from dapla_metadata.variable_definitions.generated.vardef_client.models.update_draft import (
+    UpdateDraft,
+)
 
 
 class VariableDefinition(CompleteResponse):
@@ -59,17 +61,16 @@ def get_patch(self, patch_id: int) -> CompleteResponse:
             patch_id=patch_id,
         )
 
-    @staticmethod
-    def get_id_from_short_name(short_name: str) -> str:
-        """Return id for variable definition based on short name.
-
-        Args:
-            short_name (str): The short name for one saved variable definition
-        Returns:
-            str: The id of the saved variable definition
-        """
-        variable = VariableDefinitionsApi(
+    @vardef_exception_handler
+    def update_draft(
+        self,
+        update_draft: UpdateDraft,
+    ) -> CompleteResponse:
+        """Update a Draft variable definition."""
+        return DraftVariableDefinitionsApi(
             VardefClient.get_client(),
-        ).list_variable_definitions(date_of_validity=VARDEF_EXAMPLE_DATE)
-        var_id = [var.id for var in variable if var.short_name == short_name]
-        return var_id[0]
+        ).update_variable_definition_by_id(
+            active_group="dapla-felles-developers",
+            variable_definition_id=self.id,
+            update_draft=update_draft,
+        )
diff --git a/tests/variable_definitions/test_vardef.py b/tests/variable_definitions/test_vardef.py
@@ -123,16 +123,7 @@ def test_update_draft(
 ):
     monkeypatch.setenv(DAPLA_GROUP_CONTEXT, VARDEF_EXAMPLE_ACTIVE_GROUP)
     VardefClient.set_config(client_configuration)
-    updated_draft = Vardef.update_draft(VARDEF_EXAMPLE_DEFINITION_ID, update_draft)
-    assert isinstance(updated_draft, CompleteResponse)
-
-
-def test_update_draft_by_short_name(
-    monkeypatch: pytest.MonkeyPatch,
-    client_configuration: Configuration,
-    update_draft: UpdateDraft,
-):
-    monkeypatch.setenv(DAPLA_GROUP_CONTEXT, VARDEF_EXAMPLE_ACTIVE_GROUP)
-    VardefClient.set_config(client_configuration)
-    updated_draft = Vardef.update_draft_by_short_name("landbak", update_draft)
-    assert isinstance(updated_draft, CompleteResponse)
+    landbak = Vardef.get_variable_definition(
+        variable_definition_id=VARDEF_EXAMPLE_DEFINITION_ID,
+    )
+    assert isinstance(landbak.update_draft(update_draft), CompleteResponse)