diff --git a/src/platform_reports/config.py b/src/platform_reports/config.py index bccdb9f..77a5aee 100644 --- a/src/platform_reports/config.py +++ b/src/platform_reports/config.py @@ -21,10 +21,12 @@ class Key(str): ... NEURO_PRESET_KEY = Key("platform.neuromation.io/preset") NEURO_ORG_KEY = Key("platform.neuromation.io/org") NEURO_PROJECT_KEY = Key("platform.neuromation.io/project") + NEURO_USER_KEY = Key("platform.neuromation.io/user") NEURO_JOB_KEY = Key("platform.neuromation.io/job") APOLO_ORG_KEY = Key("platform.apolo.us/org") APOLO_PROJECT_KEY = Key("platform.apolo.us/project") + APOLO_USER_KEY = Key("platform.apolo.us/user") APOLO_PRESET_KEY = Key("platform.apolo.us/preset") APOLO_APP_KEY = Key("platform.apolo.us/app") diff --git a/src/platform_reports/metrics_service.py b/src/platform_reports/metrics_service.py index e2369ae..3d8fe53 100644 --- a/src/platform_reports/metrics_service.py +++ b/src/platform_reports/metrics_service.py @@ -32,6 +32,7 @@ class CreditsUsage: resource_id: str credits: Decimal org_name: str | None = None + user_name: str | None = None class PrometheusQueryFactory: @@ -151,6 +152,12 @@ def project_name(self) -> str | None: PrometheusLabel.NEURO_PROJECT_KEY ) + @property + def user_name(self) -> str | None: + return self.labels.get(PrometheusLabel.APOLO_USER_KEY) or self.labels.get( + PrometheusLabel.NEURO_USER_KEY + ) + class StorageUsedMetric(Metric): @property @@ -188,6 +195,7 @@ def _create_for_job( category_name=CategoryName.JOBS, org_name=metric.org_name, project_name=project_name, + user_name=metric.user_name, resource_id=job_id, credits=metric.values[-1].value - metric.values[0].value, ) @@ -202,6 +210,7 @@ def _create_for_app( category_name=CategoryName.APPS, org_name=metric.org_name, project_name=project_name, + user_name=metric.user_name, resource_id=app_id, credits=metric.values[-1].value - metric.values[0].value, ) diff --git a/src/platform_reports/schema.py b/src/platform_reports/schema.py index 50e3215..b4f8347 100644 --- a/src/platform_reports/schema.py +++ b/src/platform_reports/schema.py @@ -50,8 +50,9 @@ def make_object( class PostCreditsUsageResponseSchema(Schema): - category_name = fields.Enum(CategoryName, by_value=True) + category_name = fields.Enum(CategoryName, by_value=True, required=True) org_name = fields.String() - project_name = fields.String() + project_name = fields.String(required=True) + user_name = fields.String() resource_id = fields.String(required=True) credits = fields.Decimal(required=True, as_string=True) diff --git a/tests/integration/test_api.py b/tests/integration/test_api.py index f72f7a9..3e05125 100644 --- a/tests/integration/test_api.py +++ b/tests/integration/test_api.py @@ -431,6 +431,7 @@ async def test_post_credits_usage__mocked( CreditsUsage( category_name=CategoryName.JOBS, project_name="test-project", + user_name="test-user", resource_id="test-job", credits=Decimal(1), ) @@ -457,6 +458,7 @@ async def test_post_credits_usage__mocked( "category_name": "jobs", "org_name": None, "project_name": "test-project", + "user_name": "test-user", "resource_id": "test-job", "credits": "1", } diff --git a/tests/integration/test_metrics_service.py b/tests/integration/test_metrics_service.py index c98dfe0..ea3fd43 100644 --- a/tests/integration/test_metrics_service.py +++ b/tests/integration/test_metrics_service.py @@ -126,6 +126,7 @@ async def post_query_range( "metric": { "label_platform_neuromation_io_org": "test-org", "label_platform_neuromation_io_project": "test-project", # noqa: E501 + "label_platform_neuromation_io_user": "test-user", # noqa: E501 "label_platform_neuromation_io_job": "test-job", }, "values": [[1719075883, "1"], [1719075898, "2"]], @@ -169,6 +170,7 @@ async def post_query_range( CreditsUsage( category_name=CategoryName.JOBS, project_name="test-project", + user_name="test-user", resource_id="test-job", credits=Decimal("1"), org_name="test-org", diff --git a/tests/unit/test_metrics_service.py b/tests/unit/test_metrics_service.py index c60f15b..68577ab 100644 --- a/tests/unit/test_metrics_service.py +++ b/tests/unit/test_metrics_service.py @@ -52,6 +52,7 @@ def test_create_for_compute__job(self) -> None: metric = PodCreditsMetric( labels={ "label_platform_neuromation_io_project": "test-project", + "label_platform_neuromation_io_user": "test-user", "label_platform_neuromation_io_job": "test-job", }, values=[ @@ -66,6 +67,7 @@ def test_create_for_compute__job(self) -> None: assert usage == CreditsUsage( category_name=CategoryName.JOBS, project_name="test-project", + user_name="test-user", resource_id="test-job", credits=Decimal(2), ) @@ -115,6 +117,7 @@ def test_create_for_compute__app(self) -> None: labels={ "label_platform_apolo_us_org": "test-org", "label_platform_apolo_us_project": "test-project", + "label_platform_apolo_us_user": "test-user", "label_platform_apolo_us_app": "test-app", }, values=[ @@ -130,6 +133,7 @@ def test_create_for_compute__app(self) -> None: category_name=CategoryName.APPS, org_name="test-org", project_name="test-project", + user_name="test-user", resource_id="test-app", credits=Decimal(2), ) diff --git a/tests/unit/test_schema.py b/tests/unit/test_schema.py index 26d3063..d150f85 100644 --- a/tests/unit/test_schema.py +++ b/tests/unit/test_schema.py @@ -1,10 +1,16 @@ from datetime import datetime +from decimal import Decimal -from platform_reports.schema import PostCreditsUsageRequestSchema +from platform_reports.metrics_service import CreditsUsage +from platform_reports.schema import ( + CategoryName, + PostCreditsUsageRequestSchema, + PostCreditsUsageResponseSchema, +) class TestPostCreditsUsageRequestSchema: - def test__required_fields(self) -> None: + def test_validate__required_fields(self) -> None: errors = PostCreditsUsageRequestSchema().validate({}) assert errors == { @@ -12,22 +18,22 @@ def test__required_fields(self) -> None: "start_date": ["Missing data for required field."], } - def test_category_name__incorrect(self) -> None: + def test_validate_category_name__incorrect(self) -> None: errors = PostCreditsUsageRequestSchema().validate({"category_name": "unknown"}) assert errors["category_name"][0].startswith("Must be one of:") - def test_category_name__empty_org(self) -> None: + def test_validate_category_name__empty_org(self) -> None: errors = PostCreditsUsageRequestSchema().validate({"org_name": ""}) assert errors["org_name"] == ["Shorter than minimum length 1."] - def test_category_name__empty_project(self) -> None: + def test_validate_category_name__empty_project(self) -> None: errors = PostCreditsUsageRequestSchema().validate({"project_name": ""}) assert errors["project_name"] == ["Shorter than minimum length 1."] - def test_category_name__incorrect_dates(self) -> None: + def test_validate_category_name__incorrect_dates(self) -> None: start_date = datetime.now() errors = PostCreditsUsageRequestSchema().validate( { @@ -37,3 +43,45 @@ def test_category_name__incorrect_dates(self) -> None: ) assert errors["end_date"] == ["end_date must be greater than start_date"] + + +class TestPostCreditsUsageResponseSchema: + def test_dump(self) -> None: + data = PostCreditsUsageResponseSchema().dump( + CreditsUsage( + category_name=CategoryName.JOBS, + org_name="test-org", + project_name="test-project", + user_name="test-user", + resource_id="test-job", + credits=Decimal(1), + ) + ) + + assert data == { + "category_name": "jobs", + "org_name": "test-org", + "project_name": "test-project", + "user_name": "test-user", + "resource_id": "test-job", + "credits": "1", + } + + def test_dump__defaults(self) -> None: + data = PostCreditsUsageResponseSchema().dump( + CreditsUsage( + category_name=CategoryName.JOBS, + project_name="test-project", + resource_id="test-job", + credits=Decimal(1), + ) + ) + + assert data == { + "category_name": "jobs", + "org_name": None, + "project_name": "test-project", + "user_name": None, + "resource_id": "test-job", + "credits": "1", + }