From 94c044f39cb7740f7840e4de64f4fecac0332786 Mon Sep 17 00:00:00 2001 From: Dheyay Date: Wed, 21 Aug 2024 18:10:07 -0700 Subject: [PATCH] Update handling errors due to contract expiration --- .../api/tests/test_api_u_pro_token_info_v1.py | 83 ++++++++++++++++--- uaclient/api/u/pro/token_info/v1.py | 58 +++++++++++-- 2 files changed, 125 insertions(+), 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 index bb6f594181..2328f4bc6d 100644 --- 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 index 7ab56e2c4c..c61c701dd2 100644 --- 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": """