diff --git a/python/docs/sdk/api/organization.rst b/python/docs/sdk/api/organization.rst index 5c3b93ee..47136d68 100644 --- a/python/docs/sdk/api/organization.rst +++ b/python/docs/sdk/api/organization.rst @@ -42,3 +42,31 @@ Organization .. autoclass:: BillingNamedFilterValue :members: :show-inheritance: + +.. autoclass:: BillingActivitySubfilters + :members: + :show-inheritance: + +.. autoclass:: BillingWorkloadTagSubfilter + :members: + :show-inheritance: + +.. autoclass:: BillingAssistantMessageSubfilter + :members: + :show-inheritance: + +.. autoclass:: BillingActivitySubfilterValues + :members: + :show-inheritance: + +.. autoclass:: BillingWorkloadSubfilterValues + :members: + :show-inheritance: + +.. autoclass:: BillingAssistantMessageSubfilterValues + :members: + :show-inheritance: + +.. autoclass:: BillingActivityResourceNames + :members: + :show-inheritance: diff --git a/python/lightning_sdk/__init__.py b/python/lightning_sdk/__init__.py index 19f5cc3b..00032de1 100644 --- a/python/lightning_sdk/__init__.py +++ b/python/lightning_sdk/__init__.py @@ -4,9 +4,16 @@ ActivityFileFormat, BillingActivity, BillingActivityFilterValues, + BillingActivityResourceNames, + BillingActivitySubfilters, + BillingActivitySubfilterValues, + BillingAssistantMessageSubfilter, + BillingAssistantMessageSubfilterValues, BillingDailyUsage, BillingNamedFilterValue, BillingResourceUsage, + BillingWorkloadSubfilterValues, + BillingWorkloadTagSubfilter, ) from lightning_sdk.api.teamspace_api import SecretType from lightning_sdk.api.utils import FileEntry @@ -33,9 +40,16 @@ "BillingActivityCursor", "BillingActivityFilterValues", "BillingActivityFilters", + "BillingActivityResourceNames", + "BillingActivitySubfilterValues", + "BillingActivitySubfilters", + "BillingAssistantMessageSubfilter", + "BillingAssistantMessageSubfilterValues", "BillingDailyUsage", "BillingNamedFilterValue", "BillingResourceUsage", + "BillingWorkloadSubfilterValues", + "BillingWorkloadTagSubfilter", "BucketCredentials", "CloudProvider", "ConnectionType", diff --git a/python/lightning_sdk/api/__init__.py b/python/lightning_sdk/api/__init__.py index 8b5ab2fd..3ddc2d6c 100644 --- a/python/lightning_sdk/api/__init__.py +++ b/python/lightning_sdk/api/__init__.py @@ -3,10 +3,17 @@ from lightning_sdk.api.billing_api import ( BillingActivity, BillingActivityFilterValues, + BillingActivityResourceNames, + BillingActivitySubfilters, + BillingActivitySubfilterValues, BillingApi, + BillingAssistantMessageSubfilter, + BillingAssistantMessageSubfilterValues, BillingDailyUsage, BillingNamedFilterValue, BillingResourceUsage, + BillingWorkloadSubfilterValues, + BillingWorkloadTagSubfilter, ) from lightning_sdk.api.cloud_account_api import CloudAccountApi from lightning_sdk.api.k8s_api import K8sClusterApi @@ -20,10 +27,17 @@ "ApiKeyApi", "BillingActivity", "BillingActivityFilterValues", + "BillingActivityResourceNames", + "BillingActivitySubfilterValues", + "BillingActivitySubfilters", "BillingApi", + "BillingAssistantMessageSubfilter", + "BillingAssistantMessageSubfilterValues", "BillingDailyUsage", "BillingNamedFilterValue", "BillingResourceUsage", + "BillingWorkloadSubfilterValues", + "BillingWorkloadTagSubfilter", "CloudAccountApi", "K8sClusterApi", "OrgApi", diff --git a/python/lightning_sdk/api/billing_api.py b/python/lightning_sdk/api/billing_api.py index c5ca06e1..9e9cf99c 100644 --- a/python/lightning_sdk/api/billing_api.py +++ b/python/lightning_sdk/api/billing_api.py @@ -6,14 +6,18 @@ provides the higher-level, user-facing interface built on top of this one. Three ways to get billing activity, at different grains: - - :meth:`BillingApi.get_activity`: paginated, per-resource-per-day rollup rows with raw IDs - (no resolved names). Meant for programmatic/incremental consumption. + - :meth:`BillingApi.get_activity`: paginated rollup rows, one per resource over the queried + range, keyed by raw teamspace/user/cloud account IDs. Meant for programmatic/incremental + consumption. - :meth:`BillingApi.get_session_activity`: one row per session, with resolved names. A resource (e.g. a Studio or Job) can have many sessions in the queried range. Returns CSV or JSON depending on the ``format`` argument. - :meth:`BillingApi.get_resource_activity`: one row per resource that was active in the queried range, with resolved names, a more concise view than the session-level report. Returns CSV or JSON depending on the ``format`` argument. + +The values each filter accepts come from :meth:`BillingApi.get_activity_filter_values`, and +resource IDs from :meth:`BillingApi.get_activity_filter_resource_names`. """ import csv @@ -42,7 +46,7 @@ class BillingResourceUsage: user_id: ID of the user that owns/started the resource. project_id: ID of the teamspace (project) the resource belongs to. cluster_id: ID of the cloud account the resource ran on. - resource_type: Type of the resource (e.g. "Studio", "Job"). + resource_type: Type of the resource (e.g. "lightning_code" for a Studio, "job"). resource_name: Name of the resource. created_at: When the resource was created. deleted_at: When the resource was deleted, if it has been. @@ -165,11 +169,12 @@ def _from_api(cls, data: Dict[str, Any]) -> "BillingActivity": @dataclass class BillingNamedFilterValue: - """A filterable resource ID paired with its display name. + """A filterable ID paired with its display name. Attributes: - id: ID of the resource. - name: Display name of the resource. + id: The ID to filter by. + name: Display name for the ID. Falls back to the ID itself when no name could be + resolved. """ id: Optional[str] = None @@ -180,36 +185,221 @@ def _from_api(cls, data: Dict[str, Any]) -> "BillingNamedFilterValue": return cls(id=data.get("id"), name=data.get("name")) +def _named_filter_values(items: Optional[List[Dict[str, Any]]]) -> List[BillingNamedFilterValue]: + return [BillingNamedFilterValue._from_api(item) for item in items or []] + + +@dataclass +class BillingWorkloadSubfilterValues: + """Tags a workload resource type (jobs, multi-machine jobs or deployments) can be filtered by. + + Attributes: + resource_type: Display name of the resource type these tags belong to, e.g. "Job". + tags: Tags on at least one workload of this type. Pass their IDs as + :attr:`BillingWorkloadTagSubfilter.tag_ids`. + """ + + resource_type: Optional[str] = None + tags: List[BillingNamedFilterValue] = field(default_factory=list) + + @classmethod + def _from_api(cls, data: Optional[Dict[str, Any]]) -> "BillingWorkloadSubfilterValues": + data = data or {} + return cls(resource_type=data.get("resource_type"), tags=_named_filter_values(data.get("tags"))) + + +@dataclass +class BillingAssistantMessageSubfilterValues: + """API keys assistant messages can be filtered by. + + Attributes: + api_key_ids: API keys that billed at least one message. ``name`` is the key's display + name, never its secret. Pass their IDs as + :attr:`BillingAssistantMessageSubfilter.api_key_ids`. + """ + + api_key_ids: List[BillingNamedFilterValue] = field(default_factory=list) + + @classmethod + def _from_api(cls, data: Optional[Dict[str, Any]]) -> "BillingAssistantMessageSubfilterValues": + data = data or {} + return cls(api_key_ids=_named_filter_values(data.get("api_key_ids"))) + + +@dataclass +class BillingActivitySubfilterValues: + """Filter values that only apply to one resource type, one field per type. + + Attributes: + job: Tags available to filter jobs by. + multi_machine_job: Tags available to filter multi-machine jobs by. + deployment: Tags available to filter deployments by. + assistant_message: API keys available to filter assistant messages by. + """ + + job: BillingWorkloadSubfilterValues = field(default_factory=BillingWorkloadSubfilterValues) + multi_machine_job: BillingWorkloadSubfilterValues = field(default_factory=BillingWorkloadSubfilterValues) + deployment: BillingWorkloadSubfilterValues = field(default_factory=BillingWorkloadSubfilterValues) + assistant_message: BillingAssistantMessageSubfilterValues = field( + default_factory=BillingAssistantMessageSubfilterValues + ) + + @classmethod + def _from_api(cls, data: Optional[Dict[str, Any]]) -> "BillingActivitySubfilterValues": + data = data or {} + return cls( + job=BillingWorkloadSubfilterValues._from_api(data.get("job")), + multi_machine_job=BillingWorkloadSubfilterValues._from_api(data.get("multi_machine_job")), + deployment=BillingWorkloadSubfilterValues._from_api(data.get("deployment")), + assistant_message=BillingAssistantMessageSubfilterValues._from_api(data.get("assistant_message")), + ) + + @dataclass class BillingActivityFilterValues: """The set of values a billing activity query can be filtered by, for a given scope. + Resource IDs aren't included, since an org can have too many to return at once; page through + them with :meth:`BillingApi.get_activity_filter_resource_names` instead. + Attributes: project_ids: Teamspace (project) IDs with billing activity in the queried scope. - resource_ids: Resource IDs with billing activity in the queried scope, paired with - their display names. - resource_ids_truncated: Whether ``resource_ids`` was truncated due to too many results. - resource_types: Resource types with billing activity in the queried scope. user_ids: User IDs with billing activity in the queried scope. + resource_types: Resource types with billing activity in the queried scope, e.g. + "lightning_code" (Studios), "job" or "assistant_message". + cluster_ids: Cloud account IDs with billing activity in the queried scope. + subfilters: Values for the filters that only apply to one resource type. """ project_ids: List[str] = field(default_factory=list) - resource_ids: List[BillingNamedFilterValue] = field(default_factory=list) - resource_ids_truncated: bool = False - resource_types: List[str] = field(default_factory=list) user_ids: List[str] = field(default_factory=list) + resource_types: List[str] = field(default_factory=list) + cluster_ids: List[str] = field(default_factory=list) + subfilters: BillingActivitySubfilterValues = field(default_factory=BillingActivitySubfilterValues) @classmethod def _from_api(cls, data: Dict[str, Any]) -> "BillingActivityFilterValues": return cls( project_ids=list(data.get("project_ids") or []), - resource_ids=[BillingNamedFilterValue._from_api(item) for item in data.get("resource_ids") or []], - resource_ids_truncated=bool(data.get("resource_ids_truncated")), - resource_types=list(data.get("resource_types") or []), user_ids=list(data.get("user_ids") or []), + resource_types=list(data.get("resource_types") or []), + cluster_ids=list(data.get("cluster_ids") or []), + subfilters=BillingActivitySubfilterValues._from_api(data.get("subfilters")), + ) + + +@dataclass +class BillingActivityResourceNames: + """One page of resource IDs a billing activity query can be filtered by. + + Attributes: + resource_ids: Resource IDs paired with their display names. + next_page_token: Pass as ``page_token`` to fetch the next page. Empty on the last page. + """ + + resource_ids: List[BillingNamedFilterValue] = field(default_factory=list) + next_page_token: Optional[str] = None + + @classmethod + def _from_api(cls, data: Dict[str, Any]) -> "BillingActivityResourceNames": + return cls( + resource_ids=_named_filter_values(data.get("resource_ids")), + next_page_token=data.get("next_page_token") or None, ) +@dataclass +class BillingWorkloadTagSubfilter: + """Keeps only the workloads of one resource type that carry the given tags. + + Attributes: + tag_ids: Keep workloads with any of these tags. Available tags are listed in + :attr:`BillingActivityFilterValues.subfilters`. Empty means no subfilter for this + resource type. + match_all_tags: Require every tag in ``tag_ids`` rather than any. + """ + + tag_ids: List[str] = field(default_factory=list) + match_all_tags: bool = False + + +@dataclass +class BillingAssistantMessageSubfilter: + """Keeps only the assistant messages billed through the given API keys. + + Attributes: + api_key_ids: Keep messages billed through any of these API keys. Available keys are + listed in :attr:`BillingActivityFilterValues.subfilters`. Empty means no subfilter + for assistant messages. + """ + + api_key_ids: List[str] = field(default_factory=list) + + +@dataclass +class BillingActivitySubfilters: + """Filters that only apply to one resource type, one field per type. + + Setting a field keeps only the rows of its resource type that match it. With several set, a + row matching any of them is kept. Once any field is set, rows of resource types without a + subfilter are dropped. + + Attributes: + job: Applies to job rows. + multi_machine_job: Applies to multi-machine job rows. + deployment: Applies to deployment rows. + assistant_message: Applies to assistant message rows. + """ + + job: Optional[BillingWorkloadTagSubfilter] = None + multi_machine_job: Optional[BillingWorkloadTagSubfilter] = None + deployment: Optional[BillingWorkloadTagSubfilter] = None + assistant_message: Optional[BillingAssistantMessageSubfilter] = None + + def _workload_tag_subfilters(self) -> Dict[str, BillingWorkloadTagSubfilter]: + """The workload tag subfilters that are set, keyed by their field name.""" + workloads = {"job": self.job, "multi_machine_job": self.multi_machine_job, "deployment": self.deployment} + return {name: subfilter for name, subfilter in workloads.items() if subfilter is not None and subfilter.tag_ids} + + def _api_key_ids(self) -> List[str]: + return self.assistant_message.api_key_ids if self.assistant_message is not None else [] + + +# grpc-gateway reads nested request fields from dotted, lowerCamelCase query keys. +_SUBFILTER_QUERY_KEYS = {"job": "job", "multi_machine_job": "multiMachineJob", "deployment": "deployment"} + + +def _subfilter_client_kwargs(subfilters: Optional[BillingActivitySubfilters]) -> Dict[str, Any]: + """Flatten subfilters into the kwargs the generated client takes for them.""" + if subfilters is None: + return {} + + kwargs: Dict[str, Any] = {} + for name, subfilter in subfilters._workload_tag_subfilters().items(): + kwargs[f"subfilters_{name}_tag_ids"] = subfilter.tag_ids + if subfilter.match_all_tags: + kwargs[f"subfilters_{name}_match_all_tags"] = True + if api_key_ids := subfilters._api_key_ids(): + kwargs["subfilters_assistant_message_api_key_ids"] = api_key_ids + return kwargs + + +def _subfilter_query_params(subfilters: Optional[BillingActivitySubfilters]) -> Dict[str, Any]: + """Build the subfilter query params the CSV download endpoints read.""" + if subfilters is None: + return {} + + query_params: Dict[str, Any] = {} + for name, subfilter in subfilters._workload_tag_subfilters().items(): + prefix = f"subfilters.{_SUBFILTER_QUERY_KEYS[name]}" + query_params[f"{prefix}.tagIds"] = subfilter.tag_ids + if subfilter.match_all_tags: + query_params[f"{prefix}.matchAllTags"] = "true" + if api_key_ids := subfilters._api_key_ids(): + query_params["subfilters.assistantMessage.apiKeyIds"] = api_key_ids + return query_params + + class ActivityFileFormat(str, Enum): """File format to return a billing activity report in.""" @@ -231,14 +421,16 @@ def _build_activity_query_params( resource_types: Optional[List[str]] = None, resource_ids: Optional[List[str]] = None, user_ids: Optional[List[str]] = None, + cluster_ids: Optional[List[str]] = None, + subfilters: Optional[BillingActivitySubfilters] = None, start: Optional[datetime] = None, end: Optional[datetime] = None, - limit: Optional[int] = None, - search_after: Optional[datetime] = None, - search_after_resource_id: Optional[str] = None, - search_after_resource_type: Optional[str] = None, ) -> Dict[str, Any]: - """Build the query params shared by the V2 usage report and its CSV download endpoints.""" + """Build the query params the CSV download endpoints read. + + These endpoints return the whole report at once, so unlike the V2 usage report they take no + ``limit`` or pagination cursor. + """ query_params: Dict[str, Any] = {"orgId": org_id} if project_ids is not None: query_params["projectIds"] = project_ids @@ -248,18 +440,13 @@ def _build_activity_query_params( query_params["resourceIds"] = resource_ids if user_ids is not None: query_params["userIds"] = user_ids + if cluster_ids is not None: + query_params["clusterIds"] = cluster_ids + query_params.update(_subfilter_query_params(subfilters)) if start is not None: query_params["from"] = start.isoformat() if end is not None: query_params["to"] = end.isoformat() - if limit is not None: - query_params["limit"] = limit - if search_after is not None: - query_params["searchAfter"] = search_after.isoformat() - if search_after_resource_id is not None: - query_params["searchAfterResourceId"] = search_after_resource_id - if search_after_resource_type is not None: - query_params["searchAfterResourceType"] = search_after_resource_type return query_params @@ -267,10 +454,11 @@ class BillingApi: """Internal API client for billing/usage-report requests. Combines calls through the generated OpenAPI client (:meth:`get_activity`, - :meth:`get_activity_filter_values`) with raw authenticated HTTP requests for the CSV - download endpoints (:meth:`get_session_activity`, :meth:`get_resource_activity`), which - aren't exposed by the generated client. Both return CSV or JSON depending on the ``format`` - argument; the JSON variant is built on top of the CSV download, converting it in memory. + :meth:`get_activity_filter_values`, :meth:`get_activity_filter_resource_names`) with raw + authenticated HTTP requests for the CSV download endpoints (:meth:`get_session_activity`, + :meth:`get_resource_activity`), which aren't exposed by the generated client. Both return CSV + or JSON depending on the ``format`` argument; the JSON variant is built on top of the CSV + download, converting it in memory. """ def __init__(self) -> None: @@ -284,12 +472,10 @@ def _download_activity_csv( resource_types: Optional[List[str]] = None, resource_ids: Optional[List[str]] = None, user_ids: Optional[List[str]] = None, + cluster_ids: Optional[List[str]] = None, + subfilters: Optional[BillingActivitySubfilters] = None, start: Optional[datetime] = None, end: Optional[datetime] = None, - limit: Optional[int] = None, - search_after: Optional[datetime] = None, - search_after_resource_id: Optional[str] = None, - search_after_resource_type: Optional[str] = None, ) -> str: """Download a billing activity report and return its raw CSV text.""" query_params = _build_activity_query_params( @@ -298,12 +484,10 @@ def _download_activity_csv( resource_types=resource_types, resource_ids=resource_ids, user_ids=user_ids, + cluster_ids=cluster_ids, + subfilters=subfilters, start=start, end=end, - limit=limit, - search_after=search_after, - search_after_resource_id=search_after_resource_id, - search_after_resource_type=search_after_resource_type, ) r = requests.get( @@ -328,12 +512,10 @@ def _get_activity_report( resource_types: Optional[List[str]] = None, resource_ids: Optional[List[str]] = None, user_ids: Optional[List[str]] = None, + cluster_ids: Optional[List[str]] = None, + subfilters: Optional[BillingActivitySubfilters] = None, start: Optional[datetime] = None, end: Optional[datetime] = None, - limit: Optional[int] = None, - search_after: Optional[datetime] = None, - search_after_resource_id: Optional[str] = None, - search_after_resource_type: Optional[str] = None, ) -> None: """Get a CSV-download-backed billing activity report, as CSV or JSON. @@ -355,12 +537,10 @@ def _get_activity_report( resource_types=resource_types, resource_ids=resource_ids, user_ids=user_ids, + cluster_ids=cluster_ids, + subfilters=subfilters, start=start, end=end, - limit=limit, - search_after=search_after, - search_after_resource_id=search_after_resource_id, - search_after_resource_type=search_after_resource_type, ) if format is ActivityFileFormat.CSV: @@ -379,12 +559,10 @@ def get_session_activity( resource_types: Optional[List[str]] = None, resource_ids: Optional[List[str]] = None, user_ids: Optional[List[str]] = None, + cluster_ids: Optional[List[str]] = None, + subfilters: Optional[BillingActivitySubfilters] = None, start: Optional[datetime] = None, end: Optional[datetime] = None, - limit: Optional[int] = None, - search_after: Optional[datetime] = None, - search_after_resource_id: Optional[str] = None, - search_after_resource_type: Optional[str] = None, ) -> None: """Get the session-level billing activity report, as CSV or JSON. @@ -392,10 +570,11 @@ def get_session_activity( the queried time range, so this report is the finer-grained of the two; use :meth:`get_resource_activity` for one row per resource instead. - Hits ``GET /v1/billing/usage-report/download/detailed`` with the same request shape as - the V2 usage report (see :meth:`get_activity`). Not exposed by the generated OpenAPI - client, so this issues a raw authenticated HTTP request directly against ``self._client``'s - configured host (see ``StudioApi.download_file`` for the established pattern of doing so). + Hits ``GET /v1/billing/usage-report/download/detailed`` with the same filters as the V2 + usage report (see :meth:`get_activity`), and returns the whole report at once rather than + paginating. Not exposed by the generated OpenAPI client, so this issues a raw authenticated + HTTP request directly against ``self._client``'s configured host (see + ``StudioApi.download_file`` for the established pattern of doing so). If ``format`` is :attr:`ActivityFileFormat.CSV`, the raw CSV text is written straight to ``writer``. If ``format`` is :attr:`ActivityFileFormat.JSON` (the default), the CSV is @@ -415,16 +594,12 @@ def get_session_activity( resource_ids: Restrict to these specific resource IDs. If omitted, all matching resources are returned. user_ids: Restrict to activity generated by these users. + cluster_ids: Restrict to activity on these cloud accounts. + subfilters: Restrict individual resource types further, by tag or API key. start: Only include activity on or after this time. Defaults to project/resource creation time. end: Only include activity on or before this time. Defaults to resource deletion time or now. - limit: Maximum number of entries to return. - search_after: Pagination cursor, only include entries strictly after this time. - search_after_resource_id: Pagination cursor, resource ID to break ties with - ``search_after``. Required alongside ``search_after_resource_type``. - search_after_resource_type: Pagination cursor, resource type to break ties with - ``search_after``. Required alongside ``search_after_resource_id``. """ self._get_activity_report( "/v1/billing/usage-report/download/detailed", @@ -435,12 +610,10 @@ def get_session_activity( resource_types=resource_types, resource_ids=resource_ids, user_ids=user_ids, + cluster_ids=cluster_ids, + subfilters=subfilters, start=start, end=end, - limit=limit, - search_after=search_after, - search_after_resource_id=search_after_resource_id, - search_after_resource_type=search_after_resource_type, ) def get_resource_activity( @@ -452,12 +625,10 @@ def get_resource_activity( resource_types: Optional[List[str]] = None, resource_ids: Optional[List[str]] = None, user_ids: Optional[List[str]] = None, + cluster_ids: Optional[List[str]] = None, + subfilters: Optional[BillingActivitySubfilters] = None, start: Optional[datetime] = None, end: Optional[datetime] = None, - limit: Optional[int] = None, - search_after: Optional[datetime] = None, - search_after_resource_id: Optional[str] = None, - search_after_resource_type: Optional[str] = None, ) -> None: """Get the resource-level billing activity report, as CSV or JSON. @@ -465,10 +636,11 @@ def get_resource_activity( rather than one row per session; use :meth:`get_session_activity` for the finer-grained, per-session breakdown of a resource's activity. - Hits ``GET /v1/billing/usage-report/download/summary`` with the same request shape as - the V2 usage report (see :meth:`get_activity`). Not exposed by the generated OpenAPI - client, so this issues a raw authenticated HTTP request directly against ``self._client``'s - configured host (see ``StudioApi.download_file`` for the established pattern of doing so). + Hits ``GET /v1/billing/usage-report/download/summary`` with the same filters as the V2 + usage report (see :meth:`get_activity`), and returns the whole report at once rather than + paginating. Not exposed by the generated OpenAPI client, so this issues a raw authenticated + HTTP request directly against ``self._client``'s configured host (see + ``StudioApi.download_file`` for the established pattern of doing so). If ``format`` is :attr:`ActivityFileFormat.CSV`, the raw CSV text is written straight to ``writer``. If ``format`` is :attr:`ActivityFileFormat.JSON` (the default), the CSV is @@ -488,16 +660,12 @@ def get_resource_activity( resource_ids: Restrict to these specific resource IDs. If omitted, all matching resources are returned. user_ids: Restrict to activity generated by these users. + cluster_ids: Restrict to activity on these cloud accounts. + subfilters: Restrict individual resource types further, by tag or API key. start: Only include activity on or after this time. Defaults to project/resource creation time. end: Only include activity on or before this time. Defaults to resource deletion time or now. - limit: Maximum number of entries to return. - search_after: Pagination cursor, only include entries strictly after this time. - search_after_resource_id: Pagination cursor, resource ID to break ties with - ``search_after``. Required alongside ``search_after_resource_type``. - search_after_resource_type: Pagination cursor, resource type to break ties with - ``search_after``. Required alongside ``search_after_resource_id``. """ self._get_activity_report( "/v1/billing/usage-report/download/summary", @@ -508,12 +676,10 @@ def get_resource_activity( resource_types=resource_types, resource_ids=resource_ids, user_ids=user_ids, + cluster_ids=cluster_ids, + subfilters=subfilters, start=start, end=end, - limit=limit, - search_after=search_after, - search_after_resource_id=search_after_resource_id, - search_after_resource_type=search_after_resource_type, ) def get_activity( @@ -523,6 +689,8 @@ def get_activity( resource_types: Optional[List[str]] = None, resource_ids: Optional[List[str]] = None, user_ids: Optional[List[str]] = None, + cluster_ids: Optional[List[str]] = None, + subfilters: Optional[BillingActivitySubfilters] = None, start: Optional[datetime] = None, end: Optional[datetime] = None, limit: Optional[int] = None, @@ -532,10 +700,10 @@ def get_activity( ) -> BillingActivity: """Get billing activity for an organization, optionally scoped to specific teamspaces. - Returns paginated, per-resource-per-day rollup rows keyed by raw IDs (no resolved - names); backed by the daily rollup billing activity endpoint. For a resolved-name, non-paginated - report meant for reading or exporting, use :meth:`get_session_activity` (one row - per session) or :meth:`get_resource_activity` (one row per resource) instead. + Returns paginated rollup rows, one per resource over the queried range, keyed by raw + teamspace/user/cloud account IDs. For a resolved-name, non-paginated report meant for + reading or exporting, use :meth:`get_session_activity` (one row per session) or + :meth:`get_resource_activity` (one row per resource) instead. Backed by ``billing_service_get_usage_report_v2`` (daily rollup billing activity endpoint). @@ -548,6 +716,8 @@ def get_activity( resource_ids: Restrict to these specific resource IDs. If omitted, all matching resources are returned. user_ids: Restrict to activity generated by these users. + cluster_ids: Restrict to activity on these cloud accounts. + subfilters: Restrict individual resource types further, by tag or API key. start: Only include activity on or after this time. Defaults to project/resource creation time. end: Only include activity on or before this time. Defaults to resource deletion @@ -571,6 +741,9 @@ def get_activity( kwargs["resource_ids"] = resource_ids if user_ids is not None: kwargs["user_ids"] = user_ids + if cluster_ids is not None: + kwargs["cluster_ids"] = cluster_ids + kwargs.update(_subfilter_client_kwargs(subfilters)) if start is not None: kwargs["_from"] = start if end is not None: @@ -595,8 +768,9 @@ def get_activity_filter_values( """Get the set of values a billing activity query can be filtered by. Backed by ``billing_service_get_activity_filter_values``. Returns the projects, users, - resource types, and resource IDs available to filter the activity page by, for the given - scope. + resource types, cloud accounts and per-resource-type subfilter values available to filter + the activity page by, for the given scope. Resource IDs come from + :meth:`get_activity_filter_resource_names`. Args: org_id: ID of the organization to query. @@ -612,3 +786,39 @@ def get_activity_filter_values( response = self._client.billing_service_get_activity_filter_values(**kwargs) return BillingActivityFilterValues._from_api(response.to_dict()) + + def get_activity_filter_resource_names( + self, + org_id: str, + project_id: Optional[str] = None, + search_query: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + ) -> BillingActivityResourceNames: + """Get one page of the resource IDs a billing activity query can be filtered by. + + Backed by ``billing_service_get_activity_filter_resource_names``. Only covers Studios, + jobs, multi-machine jobs and deployments, ordered by resource ID. + + Args: + org_id: ID of the organization to query. + project_id: If given, restrict to resources in this teamspace (project). + search_query: Only return resources whose name matches this. + page_size: Maximum number of resources to return. The server defaults to 20. + page_token: ``next_page_token`` from the previous page, to continue from it. + + Returns: + BillingActivityResourceNames: The page of resource IDs. + """ + kwargs: Dict[str, Any] = {"org_id": org_id} + if project_id is not None: + kwargs["project_id"] = project_id + if search_query is not None: + kwargs["search_query"] = search_query + if page_size is not None: + kwargs["page_size"] = page_size + if page_token is not None: + kwargs["page_token"] = page_token + + response = self._client.billing_service_get_activity_filter_resource_names(**kwargs) + return BillingActivityResourceNames._from_api(response.to_dict()) diff --git a/python/lightning_sdk/organization.py b/python/lightning_sdk/organization.py index 2fddd080..eb0909d0 100644 --- a/python/lightning_sdk/organization.py +++ b/python/lightning_sdk/organization.py @@ -3,7 +3,14 @@ from typing import IO, TYPE_CHECKING, List, Optional, Union from lightning_sdk.api import OrgApi -from lightning_sdk.api.billing_api import ActivityFileFormat, BillingActivity, BillingActivityFilterValues, BillingApi +from lightning_sdk.api.billing_api import ( + ActivityFileFormat, + BillingActivity, + BillingActivityFilterValues, + BillingActivityResourceNames, + BillingActivitySubfilters, + BillingApi, +) from lightning_sdk.api.org_api import MonthlySummary from lightning_sdk.owner import Owner from lightning_sdk.utils.resolve import _resolve_org_name, _resolve_teamspace @@ -16,18 +23,28 @@ class BillingActivityFilters: """Filters for narrowing down a billing activity query. + The values each filter accepts are listed by :meth:`Organization.get_activity_filter_values`, + and resource IDs by :meth:`Organization.get_activity_filter_resource_names`. + Attributes: - resource_types: Only include activity for these resource types (e.g. "Studio", "Job"). + resource_types: Only include activity for these resource types (e.g. "lightning_code" + for Studios, "job"). resource_ids: Only include activity for these specific resource IDs. user_ids: Only include activity generated by these users. + cluster_ids: Only include activity on these cloud accounts. + subfilters: Narrow individual resource types further, by tag or API key. start: Only include activity on or after this time. end: Only include activity on or before this time. - limit: Maximum number of entries to return. + limit: Maximum number of entries to return. Only applies to + :meth:`Organization.get_activity`; the session and resource activity reports always + return every entry. """ resource_types: Optional[List[str]] = None resource_ids: Optional[List[str]] = None user_ids: Optional[List[str]] = None + cluster_ids: Optional[List[str]] = None + subfilters: Optional[BillingActivitySubfilters] = None start: Optional[datetime] = None end: Optional[datetime] = None limit: Optional[int] = None @@ -194,9 +211,9 @@ def get_activity( ) -> BillingActivity: """Get billing activity for this organization, or one or more of its teamspaces. - Returns paginated, per resource rollup rows keyed by raw IDs (no resolved - names). For a resolved-name, non-paginated report meant for reading or exporting, use - :meth:`get_session_activity` (one row per session) or + Returns paginated rollup rows, one per resource over the queried range, keyed by raw + teamspace/user/cloud account IDs. For a resolved-name, non-paginated report meant for + reading or exporting, use :meth:`get_session_activity` (one row per session) or :meth:`get_resource_activity` (one row per resource) instead. Args: @@ -220,6 +237,8 @@ def get_activity( resource_types=filters.resource_types, resource_ids=filters.resource_ids, user_ids=filters.user_ids, + cluster_ids=filters.cluster_ids, + subfilters=filters.subfilters, start=filters.start, end=filters.end, limit=filters.limit, @@ -246,6 +265,39 @@ def get_activity_filter_values( return self._billing_api.get_activity_filter_values(org_id=self.id, project_id=project_id) + def get_activity_filter_resource_names( + self, + teamspace: Optional[Union[str, "Teamspace"]] = None, + search_query: Optional[str] = None, + page_size: Optional[int] = None, + page_token: Optional[str] = None, + ) -> BillingActivityResourceNames: + """Get one page of the resource IDs a billing activity query can be filtered by. + + Only covers Studios, jobs, multi-machine jobs and deployments, ordered by resource ID. + Pass the returned ``next_page_token`` back as ``page_token`` to fetch the next page. + + Args: + teamspace: Optionally scope to a single teamspace. If omitted, covers the whole + organization. + search_query: Only return resources whose name matches this. + page_size: Maximum number of resources to return. The server defaults to 20. + page_token: ``next_page_token`` from the previous page, to continue from it. + + Returns: + BillingActivityResourceNames: The page of resource IDs. + """ + resolved_teamspaces = self._resolve_billing_teamspaces(teamspace) + project_id = resolved_teamspaces[0].id if resolved_teamspaces else None + + return self._billing_api.get_activity_filter_resource_names( + org_id=self.id, + project_id=project_id, + search_query=search_query, + page_size=page_size, + page_token=page_token, + ) + def get_session_activity( self, writer: IO[str], @@ -253,7 +305,6 @@ def get_session_activity( teamspace: Optional[Union[str, "Teamspace"]] = None, teamspaces: Optional[List[Union[str, "Teamspace"]]] = None, filters: Optional[BillingActivityFilters] = None, - cursor: Optional[BillingActivityCursor] = None, ) -> None: """Get the session-level billing activity report for this organization, as CSV or JSON. @@ -271,10 +322,8 @@ def get_session_activity( teamspaces: Multiple teamspaces to scope the report to. Provide at most one of ``teamspace``/``teamspaces``. filters: Optional filters to narrow down the report. - cursor: Optional pagination cursor to continue a previous query. """ filters = filters or BillingActivityFilters() - cursor = cursor or BillingActivityCursor() resolved_teamspaces = self._resolve_billing_teamspaces(teamspace, teamspaces) self._billing_api.get_session_activity( @@ -285,12 +334,10 @@ def get_session_activity( resource_types=filters.resource_types, resource_ids=filters.resource_ids, user_ids=filters.user_ids, + cluster_ids=filters.cluster_ids, + subfilters=filters.subfilters, start=filters.start, end=filters.end, - limit=filters.limit, - search_after=cursor.search_after, - search_after_resource_id=cursor.search_after_resource_id, - search_after_resource_type=cursor.search_after_resource_type, ) def get_resource_activity( @@ -300,7 +347,6 @@ def get_resource_activity( teamspace: Optional[Union[str, "Teamspace"]] = None, teamspaces: Optional[List[Union[str, "Teamspace"]]] = None, filters: Optional[BillingActivityFilters] = None, - cursor: Optional[BillingActivityCursor] = None, ) -> None: """Get the resource-level billing activity report for this organization, as CSV or JSON. @@ -318,10 +364,8 @@ def get_resource_activity( teamspaces: Multiple teamspaces to scope the report to. Provide at most one of ``teamspace``/``teamspaces``. filters: Optional filters to narrow down the report. - cursor: Optional pagination cursor to continue a previous query. """ filters = filters or BillingActivityFilters() - cursor = cursor or BillingActivityCursor() resolved_teamspaces = self._resolve_billing_teamspaces(teamspace, teamspaces) self._billing_api.get_resource_activity( @@ -332,12 +376,10 @@ def get_resource_activity( resource_types=filters.resource_types, resource_ids=filters.resource_ids, user_ids=filters.user_ids, + cluster_ids=filters.cluster_ids, + subfilters=filters.subfilters, start=filters.start, end=filters.end, - limit=filters.limit, - search_after=cursor.search_after, - search_after_resource_id=cursor.search_after_resource_id, - search_after_resource_type=cursor.search_after_resource_type, ) def __repr__(self) -> str: diff --git a/python/tests/api/test_billing_api.py b/python/tests/api/test_billing_api.py index 224db25c..6736c61f 100644 --- a/python/tests/api/test_billing_api.py +++ b/python/tests/api/test_billing_api.py @@ -7,10 +7,15 @@ ActivityFileFormat, BillingActivity, BillingActivityFilterValues, + BillingActivityResourceNames, + BillingActivitySubfilters, BillingApi, + BillingAssistantMessageSubfilter, BillingDailyUsage, BillingNamedFilterValue, BillingResourceUsage, + BillingWorkloadSubfilterValues, + BillingWorkloadTagSubfilter, _build_activity_query_params, ) @@ -136,30 +141,51 @@ def test_billing_activity_filter_values_from_api(): filter_values = BillingActivityFilterValues._from_api( { "project_ids": ["proj-1", "proj-2"], - "resource_ids": [{"id": "res-1", "name": "my-resource"}], - "resource_ids_truncated": True, - "resource_types": ["Studio", "Job"], "user_ids": ["user-1"], + "resource_types": ["lightning_code", "job"], + "cluster_ids": ["cluster-1"], + "subfilters": { + "job": {"resource_type": "Job", "tags": [{"id": "tag-1", "name": "nightly"}]}, + "deployment": {"resource_type": "Deployment"}, + "assistant_message": {"api_key_ids": [{"id": "key-1", "name": "ci key"}]}, + }, } ) assert filter_values.project_ids == ["proj-1", "proj-2"] - assert len(filter_values.resource_ids) == 1 - assert isinstance(filter_values.resource_ids[0], BillingNamedFilterValue) - assert filter_values.resource_ids[0].id == "res-1" - assert filter_values.resource_ids_truncated is True - assert filter_values.resource_types == ["Studio", "Job"] assert filter_values.user_ids == ["user-1"] + assert filter_values.resource_types == ["lightning_code", "job"] + assert filter_values.cluster_ids == ["cluster-1"] + subfilters = filter_values.subfilters + assert subfilters.job == BillingWorkloadSubfilterValues( + resource_type="Job", tags=[BillingNamedFilterValue(id="tag-1", name="nightly")] + ) + assert subfilters.deployment == BillingWorkloadSubfilterValues(resource_type="Deployment") + assert subfilters.multi_machine_job == BillingWorkloadSubfilterValues() + assert subfilters.assistant_message.api_key_ids == [BillingNamedFilterValue(id="key-1", name="ci key")] def test_billing_activity_filter_values_from_api_missing_lists_default_to_empty(): filter_values = BillingActivityFilterValues._from_api({}) - assert filter_values.project_ids == [] - assert filter_values.resource_ids == [] - assert filter_values.resource_ids_truncated is False - assert filter_values.resource_types == [] - assert filter_values.user_ids == [] + assert filter_values == BillingActivityFilterValues() + assert filter_values.subfilters.job.tags == [] + assert filter_values.subfilters.assistant_message.api_key_ids == [] + + +def test_billing_activity_resource_names_from_api(): + page = BillingActivityResourceNames._from_api( + {"resource_ids": [{"id": "res-1", "name": "my-studio"}], "next_page_token": "tok"} + ) + + assert page.resource_ids == [BillingNamedFilterValue(id="res-1", name="my-studio")] + assert page.next_page_token == "tok" + + +def test_billing_activity_resource_names_from_api_last_page(): + page = BillingActivityResourceNames._from_api({"resource_ids": [], "next_page_token": ""}) + + assert page == BillingActivityResourceNames() # ---- _build_activity_query_params ------------------------------------------ @@ -174,34 +200,53 @@ def test_build_activity_query_params_minimal(): def test_build_activity_query_params_full(): start = datetime(2026, 1, 1, tzinfo=timezone.utc) end = datetime(2026, 2, 1, tzinfo=timezone.utc) - search_after = datetime(2026, 1, 15, tzinfo=timezone.utc) params = _build_activity_query_params( org_id="org-1", project_ids=["proj-1"], - resource_types=["Studio"], + resource_types=["assistant_message"], resource_ids=["res-1"], user_ids=["user-1"], + cluster_ids=["cluster-1"], + subfilters=BillingActivitySubfilters(assistant_message=BillingAssistantMessageSubfilter(api_key_ids=["key-1"])), start=start, end=end, - limit=10, - search_after=search_after, - search_after_resource_id="res-1", - search_after_resource_type="Studio", ) assert params == { "orgId": "org-1", "projectIds": ["proj-1"], - "resourceTypes": ["Studio"], + "resourceTypes": ["assistant_message"], "resourceIds": ["res-1"], "userIds": ["user-1"], + "clusterIds": ["cluster-1"], + "subfilters.assistantMessage.apiKeyIds": ["key-1"], "from": start.isoformat(), "to": end.isoformat(), - "limit": 10, - "searchAfter": search_after.isoformat(), - "searchAfterResourceId": "res-1", - "searchAfterResourceType": "Studio", + } + + +def test_build_activity_query_params_skips_empty_subfilters(): + subfilters = BillingActivitySubfilters( + job=BillingWorkloadTagSubfilter(), assistant_message=BillingAssistantMessageSubfilter() + ) + + assert _build_activity_query_params(org_id="org-1", subfilters=subfilters) == {"orgId": "org-1"} + + +def test_build_activity_query_params_forwards_tag_subfilters(): + subfilters = BillingActivitySubfilters( + job=BillingWorkloadTagSubfilter(tag_ids=["tag-1", "tag-2"], match_all_tags=True), + multi_machine_job=BillingWorkloadTagSubfilter(tag_ids=["tag-3"]), + deployment=BillingWorkloadTagSubfilter(tag_ids=["tag-4"]), + ) + + assert _build_activity_query_params(org_id="org-1", subfilters=subfilters) == { + "orgId": "org-1", + "subfilters.job.tagIds": ["tag-1", "tag-2"], + "subfilters.job.matchAllTags": "true", + "subfilters.multiMachineJob.tagIds": ["tag-3"], + "subfilters.deployment.tagIds": ["tag-4"], } @@ -234,6 +279,13 @@ def test_get_activity_forwards_all_kwargs(mock_client): resource_types=["Studio"], resource_ids=["res-1"], user_ids=["user-1"], + cluster_ids=["cluster-1"], + subfilters=BillingActivitySubfilters( + job=BillingWorkloadTagSubfilter(tag_ids=["tag-1", "tag-2"], match_all_tags=True), + multi_machine_job=BillingWorkloadTagSubfilter(tag_ids=["tag-3"]), + deployment=BillingWorkloadTagSubfilter(), + assistant_message=BillingAssistantMessageSubfilter(api_key_ids=["key-1"]), + ), start=start, end=end, limit=10, @@ -249,6 +301,11 @@ def test_get_activity_forwards_all_kwargs(mock_client): "resource_types": ["Studio"], "resource_ids": ["res-1"], "user_ids": ["user-1"], + "cluster_ids": ["cluster-1"], + "subfilters_job_tag_ids": ["tag-1", "tag-2"], + "subfilters_job_match_all_tags": True, + "subfilters_multi_machine_job_tag_ids": ["tag-3"], + "subfilters_assistant_message_api_key_ids": ["key-1"], "_from": start, "to": end, "limit": 10, @@ -284,6 +341,40 @@ def test_get_activity_filter_values_with_project_id(mock_client): assert call_kwargs == {"org_id": "org-1", "project_id": "proj-1"} +# ---- BillingApi.get_activity_filter_resource_names --------------------------- + + +@mock.patch("lightning_sdk.api.utils.LightningClient") +def test_get_activity_filter_resource_names_minimal(mock_client): + mock_client().billing_service_get_activity_filter_resource_names.return_value.to_dict.return_value = {} + + billing_api = BillingApi() + result = billing_api.get_activity_filter_resource_names(org_id="org-1") + + assert isinstance(result, BillingActivityResourceNames) + call_kwargs = mock_client().billing_service_get_activity_filter_resource_names.call_args[1] + assert call_kwargs == {"org_id": "org-1"} + + +@mock.patch("lightning_sdk.api.utils.LightningClient") +def test_get_activity_filter_resource_names_forwards_all_kwargs(mock_client): + mock_client().billing_service_get_activity_filter_resource_names.return_value.to_dict.return_value = {} + + billing_api = BillingApi() + billing_api.get_activity_filter_resource_names( + org_id="org-1", project_id="proj-1", search_query="train", page_size=50, page_token="tok" + ) + + call_kwargs = mock_client().billing_service_get_activity_filter_resource_names.call_args[1] + assert call_kwargs == { + "org_id": "org-1", + "project_id": "proj-1", + "search_query": "train", + "page_size": 50, + "page_token": "tok", + } + + # ---- CSV downloads ---------------------------------------------------------- @@ -360,7 +451,7 @@ def test_get_resource_activity_csv(mock_authenticate, mock_requests_get, tmp_pat format=ActivityFileFormat.CSV, writer=writer, project_ids=["proj-1"], - limit=5, + cluster_ids=["cluster-1"], ) assert result is None @@ -368,7 +459,7 @@ def test_get_resource_activity_csv(mock_authenticate, mock_requests_get, tmp_pat mock_authenticate.assert_called_once_with() call_args = mock_requests_get.call_args assert call_args[0][0].endswith("/v1/billing/usage-report/download/summary") - assert call_args[1]["params"] == {"orgId": "org-1", "projectIds": ["proj-1"], "limit": 5} + assert call_args[1]["params"] == {"orgId": "org-1", "projectIds": ["proj-1"], "clusterIds": ["cluster-1"]} @mock.patch("requests.get", autospec=True) diff --git a/python/tests/core/test_billing.py b/python/tests/core/test_billing.py index 9cca56bc..56fde557 100644 --- a/python/tests/core/test_billing.py +++ b/python/tests/core/test_billing.py @@ -4,7 +4,11 @@ import pytest -from lightning_sdk.api.billing_api import ActivityFileFormat +from lightning_sdk.api.billing_api import ( + ActivityFileFormat, + BillingActivitySubfilters, + BillingWorkloadTagSubfilter, +) from lightning_sdk.organization import BillingActivityCursor, BillingActivityFilters, Organization from lightning_sdk.teamspace import Teamspace @@ -17,6 +21,7 @@ def _make_org(org_id="org-1"): org._resolve_billing_teamspaces = Organization._resolve_billing_teamspaces.__get__(org) org.get_activity = Organization.get_activity.__get__(org) org.get_activity_filter_values = Organization.get_activity_filter_values.__get__(org) + org.get_activity_filter_resource_names = Organization.get_activity_filter_resource_names.__get__(org) org.get_session_activity = Organization.get_session_activity.__get__(org) org.get_resource_activity = Organization.get_resource_activity.__get__(org) return org @@ -90,6 +95,8 @@ def test_get_activity_defaults(): resource_types=None, resource_ids=None, user_ids=None, + cluster_ids=None, + subfilters=None, start=None, end=None, limit=None, @@ -105,9 +112,17 @@ def test_get_activity_with_filters_and_cursor(mock_resolve_teamspace): ts1 = _make_teamspace("ts-1", org) mock_resolve_teamspace.return_value = ts1 - filters = BillingActivityFilters(resource_types=["Studio"], resource_ids=["res-1"], user_ids=["user-1"], limit=5) + subfilters = BillingActivitySubfilters(job=BillingWorkloadTagSubfilter(tag_ids=["tag-1"])) + filters = BillingActivityFilters( + resource_types=["job"], + resource_ids=["res-1"], + user_ids=["user-1"], + cluster_ids=["cluster-1"], + subfilters=subfilters, + limit=5, + ) cursor = BillingActivityCursor( - search_after=datetime(2026, 1, 1), search_after_resource_id="res-1", search_after_resource_type="Studio" + search_after=datetime(2026, 1, 1), search_after_resource_id="res-1", search_after_resource_type="job" ) org.get_activity(teamspace=ts1, filters=filters, cursor=cursor) @@ -115,15 +130,17 @@ def test_get_activity_with_filters_and_cursor(mock_resolve_teamspace): org._billing_api.get_activity.assert_called_once_with( org_id="org-1", project_ids=["ts-1"], - resource_types=["Studio"], + resource_types=["job"], resource_ids=["res-1"], user_ids=["user-1"], + cluster_ids=["cluster-1"], + subfilters=subfilters, start=None, end=None, limit=5, search_after=datetime(2026, 1, 1), search_after_resource_id="res-1", - search_after_resource_type="Studio", + search_after_resource_type="job", ) @@ -150,6 +167,33 @@ def test_get_activity_filter_values_single_teamspace_scope(mock_resolve_teamspac org._billing_api.get_activity_filter_values.assert_called_once_with(org_id="org-1", project_id="ts-1") +# ---- get_activity_filter_resource_names --------------------------------- + + +def test_get_activity_filter_resource_names_org_scope(): + org = _make_org("org-1") + + result = org.get_activity_filter_resource_names() + + assert result is org._billing_api.get_activity_filter_resource_names.return_value + org._billing_api.get_activity_filter_resource_names.assert_called_once_with( + org_id="org-1", project_id=None, search_query=None, page_size=None, page_token=None + ) + + +@mock.patch("lightning_sdk.organization._resolve_teamspace") +def test_get_activity_filter_resource_names_forwards_args(mock_resolve_teamspace): + org = _make_org("org-1") + teamspace = _make_teamspace("ts-1", org) + mock_resolve_teamspace.return_value = teamspace + + org.get_activity_filter_resource_names(teamspace=teamspace, search_query="train", page_size=50, page_token="tok") + + org._billing_api.get_activity_filter_resource_names.assert_called_once_with( + org_id="org-1", project_id="ts-1", search_query="train", page_size=50, page_token="tok" + ) + + # ---- get_session_activity / get_resource_activity ----------------------- @@ -167,12 +211,10 @@ def test_get_session_activity_forwards_args(): resource_types=None, resource_ids=None, user_ids=None, + cluster_ids=None, + subfilters=None, start=None, end=None, - limit=None, - search_after=None, - search_after_resource_id=None, - search_after_resource_type=None, ) @@ -190,12 +232,10 @@ def test_get_session_activity_defaults_to_json(): resource_types=None, resource_ids=None, user_ids=None, + cluster_ids=None, + subfilters=None, start=None, end=None, - limit=None, - search_after=None, - search_after_resource_id=None, - search_after_resource_type=None, ) @@ -204,7 +244,7 @@ def test_get_resource_activity_forwards_args(mock_resolve_teamspace): org = _make_org("org-1") teamspace = _make_teamspace("ts-1", org) mock_resolve_teamspace.return_value = teamspace - filters = BillingActivityFilters(limit=100) + filters = BillingActivityFilters(cluster_ids=["cluster-1"], limit=100) writer = io.StringIO() org.get_resource_activity(format=ActivityFileFormat.CSV, writer=writer, teamspace=teamspace, filters=filters) @@ -217,12 +257,10 @@ def test_get_resource_activity_forwards_args(mock_resolve_teamspace): resource_types=None, resource_ids=None, user_ids=None, + cluster_ids=["cluster-1"], + subfilters=None, start=None, end=None, - limit=100, - search_after=None, - search_after_resource_id=None, - search_after_resource_type=None, ) @@ -240,10 +278,8 @@ def test_get_resource_activity_defaults_to_json(): resource_types=None, resource_ids=None, user_ids=None, + cluster_ids=None, + subfilters=None, start=None, end=None, - limit=None, - search_after=None, - search_after_resource_id=None, - search_after_resource_type=None, )