From a2b0ad8c686342b77a40e369b168418ed5bd157a Mon Sep 17 00:00:00 2001
From: Jure Zmrzlikar <zmrzlikar.jure@gmail.com>
Date: Mon, 13 Jun 2022 14:57:53 +0200
Subject: [PATCH] Support setting billing account on Collection

---
 docs/CHANGELOG.rst                |  9 +++++++++
 src/resdk/resources/collection.py | 16 ++++++++++++++++
 src/resdk/resources/utils.py      | 15 +++++++++++++++
 3 files changed, 40 insertions(+)

diff --git a/docs/CHANGELOG.rst b/docs/CHANGELOG.rst
index eb259619..a4409b42 100644
--- a/docs/CHANGELOG.rst
+++ b/docs/CHANGELOG.rst
@@ -5,6 +5,15 @@ Change Log
 All notable changes to this project are documented in this file.
 
 
+==========
+Unreleased
+==========
+
+Added
+-----
+- Support setting billing account on ``Collection`` resource
+
+
 ===================
 15.0.0 - 2022-06-06
 ===================
diff --git a/src/resdk/resources/collection.py b/src/resdk/resources/collection.py
index 309c4ca1..f1a75b3d 100644
--- a/src/resdk/resources/collection.py
+++ b/src/resdk/resources/collection.py
@@ -1,11 +1,13 @@
 """Collection resources."""
 import logging
+from urllib.parse import urljoin
 
 from resdk.shortcuts.collection import CollectionRelationsMixin
 
 from ..utils.decorators import assert_object_exists
 from .base import BaseResolweResource
 from .descriptor import DescriptorSchema
+from .utils import _get_billing_account_id
 
 
 class BaseCollection(BaseResolweResource):
@@ -203,3 +205,17 @@ def duplicate(self):
         """
         duplicated = self.api().duplicate.post({"ids": [self.id]})
         return self.__class__(resolwe=self.resolwe, **duplicated[0])
+
+    @assert_object_exists
+    def assign_to_billing_account(self, billing_account_name):
+        """Assign given collection to a billing account."""
+        billing_account_id = _get_billing_account_id(self.resolwe, billing_account_name)
+
+        # Assign collection to a billing account
+        response = self.resolwe.session.post(
+            urljoin(
+                self.resolwe.url, f"api/billingaccount/{billing_account_id}/collection"
+            ),
+            data={"collection_id": self.id},
+        )
+        response.raise_for_status()
diff --git a/src/resdk/resources/utils.py b/src/resdk/resources/utils.py
index 3376f408..4368a957 100644
--- a/src/resdk/resources/utils.py
+++ b/src/resdk/resources/utils.py
@@ -1,5 +1,7 @@
 """Resource utility functions."""
+import functools
 from datetime import datetime
+from urllib.parse import urljoin
 
 import pytz
 import tzlocal
@@ -204,3 +206,16 @@ def parse_resolwe_datetime(dtime):
         local_time = utc_aware.astimezone(local_tz)
 
         return local_time
+
+
+@functools.lru_cache(128)
+def _get_billing_account_id(res, name):
+    """Get billing account ID based on it's name."""
+    response = res.session.get(urljoin(res.url, f"api/billingaccount"))
+    if not response.json():
+        raise ValueError(f"You do not have sufficient permissions for assigning billing accounts.")
+    for item in response.json():
+        if item["name"] == name:
+            return item["id"]
+    else:
+        raise ValueError(f'Could not find a billing account with name "{name}"')