First_agent_template / tools /helm_test_gen.py
h2oaimichalmarszalek's picture
Update tools/helm_test_gen.py
c087470 verified
raw
history blame contribute delete
3.21 kB
import re
import os
import yaml
from typing import Optional
from smolagents import Tool
from jinja2 import Environment, FileSystemLoader
from pydash import py_
class GenerateHelmTest(Tool):
name = "generate_helm_tests"
description = "Performs a duckduckgo web search based on your query (think a Google search) then returns the top search results."
inputs = {'document': {'type': 'string', 'description': 'A YAML string containing one or more Kubernetes resource definitions.'}}
output_type = "string"
SEP = '''
---
'''
@staticmethod
def get_yaml_value(d: dict[str], xpath: str, default: Optional[str] = None) -> str:
try:
if '.' in xpath:
h = xpath.split('.')[0]
t = '.'.join(xpath.split('.')[1:])
return GenerateHelmTest.get_yaml_value(d[h], t)
else:
return d[xpath]
except Exception as ex:
if default:
return default
else:
raise ex
@staticmethod
def bool_to_str(value: bool) -> str:
return str(value).lower()
@staticmethod
def get_object_by_kind(document: dict[str], kinds: tuple[str]) -> list[dict[str]]:
return [*py_.filter_(document, lambda x: x['kind'] in kinds)]
@staticmethod
def get_envs_by_ref(container: dict[str], ref: str, name: str) -> list[dict[str]]:
return [*py_.filter(container['env'], lambda x: x.get('valueFrom', {f'{ref}': {}}).get(f'{ref}', {'name': ''}).get('name') == name)] # noqa: E501
@staticmethod
def get_binding_subjects(subjects: dict[str], name: str) -> list[dict[str]]:
return [*py_.filter(subjects, lambda x: x['name'] == name)]
def __init__(self, *args, **kwargs):
self.is_initialized = False
def forward(self, document: str) -> str:
"""Generate Helm test scenarios from a Kubernetes YAML document.
Args:
document (str): A YAML string containing one or more Kubernetes resource definitions.
Returns:
str: A string containing the generated test scenarios, separated by a predefined
separator when multiple resources are processed.
"""
doc = [*yaml.safe_load_all(document)]
result = ''
current_dir = os.path.dirname(os.path.abspath(__file__))
template_dir = os.path.join(current_dir, 'templates')
for item in doc:
name = item['metadata']['name']
kind = item['kind']
kindDesc = (re.sub(r'((?<=[a-z])[A-Z]|(?<!\A)[A-Z](?=[a-z]))', r' \1', kind)).lower()
env = Environment(loader=FileSystemLoader(template_dir))
env.globals['gyv'] = GenerateHelmTest.get_yaml_value
env.globals['bool_to_str'] = GenerateHelmTest.bool_to_str
env.globals['get_object_by_kind'] = GenerateHelmTest.get_object_by_kind
env.globals['get_env_by_ref'] = GenerateHelmTest.get_envs_by_ref
env.globals['get_binding_subjects'] = GenerateHelmTest.get_binding_subjects
template = env.get_template('scenario.jinja')
output = template.render(name=name, app='app', kind=kind, kindDesc=kindDesc, item=item, document=doc)
if result:
result += GenerateHelmTest.SEP
result += output
return f'```yaml\n{result}\n```'