From 9cf9e58fdc92d10b7188942678c7ccb8732359ad Mon Sep 17 00:00:00 2001
From: Anthony Mahanna <anthony.mahanna@gmail.com>
Date: Mon, 11 Sep 2023 22:20:30 -0400
Subject: [PATCH 1/2] initial commit

---
 arango/database.py     | 50 ++++++++++++++++++++++++++++++++++++++++++
 arango/exceptions.py   |  8 +++++++
 tests/test_database.py | 17 ++++++++++++++
 3 files changed, 75 insertions(+)

diff --git a/arango/database.py b/arango/database.py
index 875a04e0..2d683d6f 100644
--- a/arango/database.py
+++ b/arango/database.py
@@ -43,6 +43,8 @@
     ServerEchoError,
     ServerEncryptionError,
     ServerEngineError,
+    ServerLicenseGetError,
+    ServerLicenseSetError,
     ServerLogLevelError,
     ServerLogLevelSetError,
     ServerMetricsError,
@@ -348,6 +350,54 @@ def response_handler(resp: Response) -> Json:
 
         return self._execute(request, response_handler)
 
+    def license(self) -> Result[Json]:
+        """View the license information and status of an
+        Enterprise Edition instance. Can be called on
+        single servers, Coordinators, and DB-Servers.
+
+        :return: Server license.
+        :rtype: dict
+        :raise arango.exceptions.ServerLicenseGetError: If retrieval fails.
+        """
+        request = Request(method="get", endpoint="/_admin/license")
+
+        def response_handler(resp: Response) -> Json:
+            if resp.is_success:
+                result: Json = resp.body
+                return result
+            raise ServerLicenseGetError(resp, request)
+
+        return self._execute(request, response_handler)
+
+    def set_license(self, license: str, force: bool = False) -> Result[Json]:
+        """Set a new license for an Enterprise Edition
+        instance. Can be called on single servers, Coordinators,
+        and DB-Servers.
+
+        :param license: The Base64-encoded license string.
+        :type license: str
+        :param force: If set to True, the new license will be set even if
+            it expires sooner than the current license.
+        :type force: bool
+        :return: Server license.
+        :rtype: dict
+        :raise arango.exceptions.ServerLicenseError: If retrieval fails.
+        """
+        request = Request(
+            method="put",
+            endpoint="/_admin/license",
+            params={"force": force},
+            data=license,
+        )
+
+        def response_handler(resp: Response) -> Json:
+            if resp.is_success:
+                result: Json = resp.body
+                return result
+            raise ServerLicenseSetError(resp, request)
+
+        return self._execute(request, response_handler)
+
     def status(self) -> Result[Json]:
         """Return ArangoDB server status.
 
diff --git a/arango/exceptions.py b/arango/exceptions.py
index f9b6de2d..6fed37b8 100644
--- a/arango/exceptions.py
+++ b/arango/exceptions.py
@@ -614,6 +614,14 @@ class ServerDetailsError(ArangoServerError):
     """Failed to retrieve server details."""
 
 
+class ServerLicenseGetError(ArangoServerError):
+    """Failed to retrieve server license."""
+
+
+class ServerLicenseSetError(ArangoServerError):
+    """Failed to set server license."""
+
+
 class ServerStatusError(ArangoServerError):
     """Failed to retrieve server status."""
 
diff --git a/tests/test_database.py b/tests/test_database.py
index 60685c0f..d00c742f 100644
--- a/tests/test_database.py
+++ b/tests/test_database.py
@@ -319,3 +319,20 @@ def test_database_utf8(sys_db, db_version, special_db_names):
         assert sys_db.create_database(name)
         assert sys_db.has_database(name)
         assert sys_db.delete_database(name)
+
+
+def test_license(sys_db, db, enterprise):
+    license = sys_db.license()
+    assert isinstance(license, dict)
+
+    if enterprise:
+        assert set(license.keys()) == {
+            "upgrading",
+            "features",
+            "hash",
+            "license",
+            "version",
+            "status",
+        }
+    else:
+        assert license == {"license": "none"}

From a416b18b8ff766bde5288e8fcc30bed2b39e1bd3 Mon Sep 17 00:00:00 2001
From: Anthony Mahanna <anthony.mahanna@gmail.com>
Date: Fri, 15 Sep 2023 11:13:05 -0400
Subject: [PATCH 2/2] new: set license test for community

---
 tests/test_database.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/test_database.py b/tests/test_database.py
index d00c742f..0e5d159c 100644
--- a/tests/test_database.py
+++ b/tests/test_database.py
@@ -20,6 +20,7 @@
     ServerDetailsError,
     ServerEchoError,
     ServerEngineError,
+    ServerLicenseSetError,
     ServerLogLevelError,
     ServerLogLevelSetError,
     ServerMetricsError,
@@ -321,7 +322,7 @@ def test_database_utf8(sys_db, db_version, special_db_names):
         assert sys_db.delete_database(name)
 
 
-def test_license(sys_db, db, enterprise):
+def test_license(sys_db, enterprise):
     license = sys_db.license()
     assert isinstance(license, dict)
 
@@ -336,3 +337,5 @@ def test_license(sys_db, db, enterprise):
         }
     else:
         assert license == {"license": "none"}
+        with pytest.raises(ServerLicenseSetError):
+            sys_db.set_license("abc")