From 262251112a4d64c17b874194e075b948056d3edf Mon Sep 17 00:00:00 2001 From: Brahim Boukoufallah Date: Fri, 28 Jun 2024 16:31:34 +0200 Subject: [PATCH 1/2] feat: allow to override parquet_prefix and include idle costs --- src/opencost_parquet_exporter.py | 31 +++++++++++++++++++++------ src/test_opencost_parquet_exporter.py | 21 ++++++++++++++++++ 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/src/opencost_parquet_exporter.py b/src/opencost_parquet_exporter.py index d118783..b6f0a81 100644 --- a/src/opencost_parquet_exporter.py +++ b/src/opencost_parquet_exporter.py @@ -18,8 +18,11 @@ def get_config( window_end=None, s3_bucket=None, file_key_prefix=None, + parquet_prefix=None, aggregate_by=None, - step=None): + step=None, + include_idle=None, + idle_by_node=None): """ Get configuration for the parquet exporter based on either provided parameters or environment variables. @@ -48,6 +51,12 @@ def get_config( - step (str): Granularity for the data aggregation, defaults to the 'OPENCOST_PARQUET_STEP' environment variable, or '1h' if not set. + - include_idle (str): Include idle resources in the data, + defaults to the 'OPENCOST_PARQUET_INCLUDE_IDLE' environment variable, + or 'false' if not set. + - idle_by_node (str): Include idle resources by node, + defaults to the 'OPENCOST_PARQUET_IDLE_BY_NODE' environment variable, + or 'false' if not set. Returns: - dict: Configuration dictionary with keys for 'url', 'params', 's3_bucket', @@ -74,6 +83,10 @@ def get_config( aggregate_by = os.environ.get('OPENCOST_PARQUET_AGGREGATE', 'namespace,pod,container') if step is None: step = os.environ.get('OPENCOST_PARQUET_STEP', '1h') + if include_idle is None: + include_idle = os.environ.get('OPENCOST_PARQUET_INCLUDE_IDLE', 'false') + if idle_by_node is None: + idle_by_node = os.environ.get('OPENCOST_PARQUET_IDLE_BY_NODE', 'false') if s3_bucket is not None: config['s3_bucket'] = s3_bucket @@ -85,12 +98,20 @@ def get_config( datetime.now() - timedelta(1), '%Y-%m-%d') window_start = yesterday+'T00:00:00Z' window_end = yesterday+'T23:59:59Z' + if parquet_prefix is None: + window = datetime.strptime(window_start, "%Y-%m-%dT%H:%M:%SZ") + parquet_prefix = os.environ.get( + 'OPENCOST_PARQUET_PREFIX', + f"{file_key_prefix}/year={window.year}/month={window.month}/day={window.day}" + ) + + config['parquet_prefix'] = parquet_prefix window = f"{window_start},{window_end}" config['params'] = ( ("window", window), ("aggregate", aggregate_by), - ("includeIdle", "false"), - ("idleByNode", "false"), + ("includeIdle", include_idle), + ("idleByNode", idle_by_node), ("includeProportionalAssetResourceCosts", "false"), ("format", "json"), ("step", step) @@ -244,9 +265,7 @@ def save_result(processed_result, config): - uri : String with the path where the data was saved. """ file_name = 'k8s_opencost.parquet' - window = datetime.strptime(config['window_start'], "%Y-%m-%dT%H:%M:%SZ") - parquet_prefix = f"{config['file_key_prefix']}/year={window.year}"\ - f"/month={window.month}/day={window.day}" + parquet_prefix = config['parquet_prefix'] try: if config.get('s3_bucket', None): uri = f"s3://{config['s3_bucket']}/{parquet_prefix}/{file_name}" diff --git a/src/test_opencost_parquet_exporter.py b/src/test_opencost_parquet_exporter.py index d0bc53c..91b4f3d 100644 --- a/src/test_opencost_parquet_exporter.py +++ b/src/test_opencost_parquet_exporter.py @@ -90,6 +90,27 @@ def test_get_config_no_window_end(self): self.assertNotIn('s3_bucket', config) self.assertEqual(config['params'][0][1], window) + def test_get_config_with_parquet_prefix_env_var(self): + """Test get_config overrides parquet prefix when env var is set.""" + with patch.dict(os.environ, { + 'OPENCOST_PARQUET_INCLUDE_IDLE': 'true', + 'OPENCOST_PARQUET_IDLE_BY_NODE': 'true'}, clear=True): + config = get_config() + self.assertEqual(config['params'][2][0], 'includeIdle') + self.assertEqual(config['params'][2][1], 'true') + + self.assertEqual(config['params'][3][0], 'idleByNode') + self.assertEqual(config['params'][3][1], 'true') + + def test_get_config_with_idle_env_vars(self): + """Test get_config overrides parquet prefix when env var is set.""" + with patch.dict(os.environ, { + 'OPENCOST_PARQUET_FILE_KEY_PREFIX': 'test-prefix/', + 'OPENCOST_PARQUET_PREFIX': 'my/parquet/prefix', + 'OPENCOST_PARQUET_STEP': '1m'}, clear=True): + config = get_config() + self.assertEqual(config['parquet_prefix'], 'my/parquet/prefix') + class TestRequestData(unittest.TestCase): """ Test request_data method """ @patch('opencost_parquet_exporter.requests.get') From ee2b6a9dc6ad97545160b16c78454fad551bdc45 Mon Sep 17 00:00:00 2001 From: Brahim Boukoufallah Date: Sat, 29 Jun 2024 17:23:24 +0200 Subject: [PATCH 2/2] chore: add more test cases --- src/test_opencost_parquet_exporter.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/src/test_opencost_parquet_exporter.py b/src/test_opencost_parquet_exporter.py index 91b4f3d..8d6f012 100644 --- a/src/test_opencost_parquet_exporter.py +++ b/src/test_opencost_parquet_exporter.py @@ -2,6 +2,7 @@ import unittest from unittest.mock import patch, MagicMock import os +from datetime import datetime import requests from freezegun import freeze_time from opencost_parquet_exporter import get_config, request_data @@ -90,7 +91,7 @@ def test_get_config_no_window_end(self): self.assertNotIn('s3_bucket', config) self.assertEqual(config['params'][0][1], window) - def test_get_config_with_parquet_prefix_env_var(self): + def test_get_config_with_idle_env_vars(self): """Test get_config overrides parquet prefix when env var is set.""" with patch.dict(os.environ, { 'OPENCOST_PARQUET_INCLUDE_IDLE': 'true', @@ -102,14 +103,26 @@ def test_get_config_with_parquet_prefix_env_var(self): self.assertEqual(config['params'][3][0], 'idleByNode') self.assertEqual(config['params'][3][1], 'true') - def test_get_config_with_idle_env_vars(self): + @freeze_time('2024-06-29') + def test_get_config_with_parquet_prefix_env_var(self): + + yesterday = '2024-06-28' + window_start = yesterday+'T01:00:00Z' + window = datetime.strptime(window_start, "%Y-%m-%dT%H:%M:%SZ") + """Test get_config overrides parquet prefix when env var is set.""" with patch.dict(os.environ, { 'OPENCOST_PARQUET_FILE_KEY_PREFIX': 'test-prefix/', - 'OPENCOST_PARQUET_PREFIX': 'my/parquet/prefix', + 'OPENCOST_PARQUET_PREFIX': f"year={window.year}/month={window.month}/day={window.day}/hour={window.hour}", 'OPENCOST_PARQUET_STEP': '1m'}, clear=True): config = get_config() - self.assertEqual(config['parquet_prefix'], 'my/parquet/prefix') + self.assertEqual(config['parquet_prefix'], 'year=2024/month=6/day=28/hour=1') + + @freeze_time('2024-06-29') + def test_get_config_with_default_parquet_prefix_env_var(self): + """Test get_config with default parquet prefix when env var is not set.""" + config = get_config() + self.assertEqual(config['parquet_prefix'], '/tmp//year=2024/month=6/day=28') class TestRequestData(unittest.TestCase): """ Test request_data method """