Update handling errors due to contract expiration

canonical · Sep 24, 2024 · 94c044f · 94c044f
1 parent f663a8d
commit 94c044f
Show file tree

Hide file tree

Showing 2 changed files with 125 additions and 16 deletions.
diff --git a/uaclient/api/tests/test_api_u_pro_token_info_v1.py b/uaclient/api/tests/test_api_u_pro_token_info_v1.py
@@ -1,6 +1,9 @@
+import copy
+
 import mock
+import pytest
 
-from uaclient import util
+from uaclient import exceptions, util
 from uaclient.api.u.pro.token_info.v1 import (
     AccountInfo,
     ContractInfo,
@@ -44,16 +47,15 @@
 ]
 
 
-@mock.patch(
-    "uaclient.api.u.pro.token_info.v1.get_available_resources",
-    return_value=RESPONSE_AVAILABLE_SERVICES,
-)
-@mock.patch(
-    "uaclient.api.u.pro.token_info.v1.get_contract_information",
-    return_value=RESPONSE_CONTRACT_INFO,
-)
 class TestTokenInfo:
-
+    @mock.patch(
+        "uaclient.api.u.pro.token_info.v1.get_available_resources",
+        return_value=RESPONSE_AVAILABLE_SERVICES,
+    )
+    @mock.patch(
+        "uaclient.api.u.pro.token_info.v1.get_contract_information",
+        return_value=RESPONSE_CONTRACT_INFO,
+    )
     def test_token_info_output(
         self, m_get_contract_information, m_get_available_resources, FakeConfig
     ):
@@ -63,7 +65,12 @@ def test_token_info_output(
         )
         expected_info = TokenInfoResult(
             account=AccountInfo(id="some_id", name="Name"),
-            contract=ContractInfo(id="some_id", name="Name"),
+            contract=ContractInfo(
+                id="some_id",
+                name="Name",
+                effective=None,
+                expires=util.parse_rfc3339_date("9999-12-31T00:00:00Z"),
+            ),
             services=[
                 ServiceInfo(
                     name="livepatch",
@@ -75,3 +82,57 @@ def test_token_info_output(
             ],
         )
         assert token_info == expected_info
+
+    @pytest.mark.parametrize(
+        "expected_error,expected_err_message,contract_field,date_value",
+        (
+            (
+                exceptions.TokenForbiddenExpired,
+                (
+                    'Contract "some_id" expired on December 31, 2019\n'
+                    "Visit https://ubuntu.com/pro/dashboard to manage "
+                    "contract tokens."
+                ),
+                "effectiveTo",
+                util.parse_rfc3339_date("2019-12-31T00:00:00Z"),
+            ),
+            (
+                exceptions.TokenForbiddenNotYet,
+                (
+                    'Contract "some_id" is not effective until December 31, '
+                    "9999\n"
+                    "Visit https://ubuntu.com/pro/dashboard to manage "
+                    "contract tokens."
+                ),
+                "effectiveFrom",
+                util.parse_rfc3339_date("9999-12-31T00:00:00Z"),
+            ),
+        ),
+    )
+    @mock.patch(
+        "uaclient.api.u.pro.token_info.v1.get_available_resources",
+        return_value=RESPONSE_AVAILABLE_SERVICES,
+    )
+    @mock.patch("uaclient.api.u.pro.token_info.v1.get_contract_information")
+    def test_attach_forbidden_error(
+        self,
+        m_get_contract_information,
+        m_get_available_resources,
+        expected_error,
+        expected_err_message,
+        contract_field,
+        date_value,
+        capsys,
+        FakeConfig,
+    ):
+        exp_contract = copy.deepcopy(RESPONSE_CONTRACT_INFO)
+        exp_contract["contractInfo"][contract_field] = date_value
+        m_get_contract_information.return_value = exp_contract
+        cfg = FakeConfig()
+        try:
+            _get_token_info(
+                options=TokenInfoOptions(token="contract_token"), cfg=cfg
+            )
+        except expected_error as e:
+            err_message = str(e)
+            assert err_message == expected_err_message
diff --git a/uaclient/api/u/pro/token_info/v1.py b/uaclient/api/u/pro/token_info/v1.py
@@ -1,5 +1,6 @@
 import logging
-from typing import Any, Dict, List
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
 
 from uaclient import exceptions, util
 from uaclient.api.api import APIEndpoint
@@ -9,10 +10,12 @@
 from uaclient.data_types import (
     BoolDataValue,
     DataObject,
+    DatetimeDataValue,
     Field,
     StringDataValue,
     data_list,
 )
+from uaclient.defaults import ATTACH_FAIL_DATE_FORMAT
 from uaclient.entitlements import entitlement_factory
 
 LOG = logging.getLogger(util.replace_top_level_logger_name(__name__))
@@ -44,11 +47,21 @@ class ContractInfo(DataObject):
     fields = [
         Field("id", StringDataValue),
         Field("name", StringDataValue),
+        Field("effective", DatetimeDataValue),
+        Field("expires", DatetimeDataValue),
     ]
 
-    def __init__(self, id: str, name: str):
+    def __init__(
+        self,
+        id: str,
+        name: str,
+        effective: Optional[datetime],
+        expires: Optional[datetime],
+    ):
         self.id = id
         self.name = name
+        self.effective = effective
+        self.expires = expires
 
 
 class ServiceInfo(DataObject):
@@ -124,15 +137,42 @@ def _get_token_info(
         raise e
 
     contract_info = contract_information.get("contractInfo", {})
-    contract = ContractInfo(
-        id=contract_info.get("id", ""), name=contract_info.get("name", "")
-    )
 
     account_info = contract_information.get("accountInfo", {})
     account = AccountInfo(
         id=account_info.get("id", ""), name=account_info.get("name", "")
     )
 
+    # Check contract expiration
+    now = datetime.now(timezone.utc)
+    if contract_info.get("effectiveTo"):
+        expiration_datetime = contract_info.get("effectiveTo")
+        delta = expiration_datetime - now
+        if delta.total_seconds() <= 0:
+            raise exceptions.TokenForbiddenExpired(
+                contract_id=contract_info.get("id", ""),
+                date=expiration_datetime.strftime(ATTACH_FAIL_DATE_FORMAT),
+                contract_expiry_date=expiration_datetime.strftime("%m-%d-%Y"),
+            )
+    if contract_info.get("effectiveFrom"):
+        effective_datetime = contract_info.get("effectiveFrom")
+        delta = now - effective_datetime
+        if delta.total_seconds() <= 0:
+            raise exceptions.TokenForbiddenNotYet(
+                contract_id=contract_info.get("id", ""),
+                date=effective_datetime.strftime(ATTACH_FAIL_DATE_FORMAT),
+                contract_effective_date=effective_datetime.strftime(
+                    "%m-%d-%Y"
+                ),
+            )
+
+    contract = ContractInfo(
+        id=contract_info.get("id", ""),
+        name=contract_info.get("name", ""),
+        effective=contract_info.get("effectiveFrom", None),
+        expires=contract_info.get("effectiveTo", None),
+    )
+
     services = []
     resources = get_available_resources(cfg)
     inapplicable_resources = {
@@ -190,6 +230,14 @@ def _get_token_info(
             exceptions.AttachInvalidTokenError,
             "When an invalid token is passed as an argument",
         ),
+        (
+            exceptions.TokenForbiddenExpired,
+            "When the contract has expired",
+        ),
+        (
+            exceptions.TokenForbiddenNotYet,
+            "When the contract is not yet effective",
+        ),
     ],
     "example_cli": "pro api u.pro.token_info.v1 --args token=CONTRACT_TOKEN",
     "example_json": """