def test_project_secrets_credentials(self):
# This simulates running a job in a pod with project-secrets assigned to it
for param in credential_params:
os.environ.pop(param, None)
os.environ[SecretsStore.k8s_env_variable_name_for_secret(
param)] = config["env"][param]
self._perform_aws_s3_tests()
# cleanup
for param in credential_params:
os.environ.pop(
SecretsStore.k8s_env_variable_name_for_secret(param))
def __init__(
self,
graph=None,
parameters=None,
load_mode=None,
function_uri=None,
verbose=False,
version=None,
functions=None,
graph_initializer=None,
error_stream=None,
track_models=None,
secret_sources=None,
):
self._graph = None
self.graph: Union[RouterState, RootFlowState] = graph
self.function_uri = function_uri
self.parameters = parameters or {}
self.verbose = verbose
self.load_mode = load_mode or "sync"
self.version = version or "v2"
self.context = None
self._current_function = None
self.functions = functions or {}
self.graph_initializer = graph_initializer
self.error_stream = error_stream
self.track_models = track_models
self._error_stream_object = None
self.secret_sources = secret_sources
self._secrets = SecretsStore.from_list(secret_sources)
self._db_conn = None
self.resource_cache = None
def _get_credentials_string(self):
gcp_project = self.attributes.get("gcp_project", None)
key = "GCP_CREDENTIALS"
gcp_cred_string = os.getenv(key) or os.getenv(
SecretsStore.k8s_env_variable_name_for_secret(key)
)
return gcp_cred_string, gcp_project
def test_using_dataitem_secrets(self):
# make sure no other auth method is configured
for param in credential_params:
os.environ.pop(param, None)
os.environ.pop(
SecretsStore.k8s_env_variable_name_for_secret(param), None)
secrets = {param: config["env"][param] for param in credential_params}
self._perform_aws_s3_tests(secrets=secrets)
def test_inline_str():
spec = {
"secret_sources": [{
"kind": "inline",
"source": "{'abc': 'def'}"
}],
}
ss = SecretsStore.from_list(spec["secret_sources"])
assert ss.get("abc") == "def", "failed on 1st env var secret"
def test_inline_str():
spec = {
'secret_sources': [{
'kind': 'inline',
'source': "{'abc': 'def'}"
}],
}
ss = SecretsStore.from_list(spec['secret_sources'])
assert ss.get('abc') == 'def', 'failed on 1st env var secret'
def test_load():
environ['ENV123'] = 'xx'
environ['ENV456'] = 'yy'
ss = SecretsStore.from_list(spec['secret_sources'])
assert ss.get('ENV123') == 'xx', 'failed on 1st env var secret'
assert ss.get('ENV456') == 'yy', 'failed on 1st env var secret'
assert ss.get('MYENV') == '123', 'failed on 1st env var secret'
assert ss.get('MY2NDENV') == '456', 'failed on 1st env var secret'
assert ss.get('abc') == 'def', 'failed on 1st env var secret'
print(ss.get_all())
def test_load():
environ["ENV123"] = "xx"
environ["ENV456"] = "yy"
ss = SecretsStore.from_list(spec["secret_sources"])
assert ss.get("ENV123") == "xx", "failed on 1st env var secret"
assert ss.get("ENV456") == "yy", "failed on 1st env var secret"
assert ss.get("MYENV") == "123", "failed on 1st env var secret"
assert ss.get("MY2NDENV") == "456", "failed on 1st env var secret"
assert ss.get("abc") == "def", "failed on 1st env var secret"
print(ss.items())
def test_using_env_variables(self):
# Use "naked" env variables, useful in client-side sdk.
for param in credential_params:
os.environ[param] = config["env"][param]
os.environ.pop(
SecretsStore.k8s_env_variable_name_for_secret(param), None)
self._perform_aws_s3_tests()
# cleanup
for param in credential_params:
os.environ.pop(param)
def get_expected_env_variables_from_secrets(self,
project,
encode_key_names=True,
include_internal=False):
expected_env_from_secrets = {}
secret_name = mlrun.api.utils.singletons.k8s.get_k8s(
).get_project_secret_name(project)
for key in self.project_secrets_map.get(project, {}):
if key.startswith("mlrun.") and not include_internal:
continue
env_variable_name = (
SecretsStore.k8s_env_variable_name_for_secret(key)
if encode_key_names else key)
expected_env_from_secrets[env_variable_name] = {secret_name: key}
return expected_env_from_secrets
def init(
self,
context,
namespace,
resource_cache: ResourceCache = None,
logger=None,
is_mock=False,
):
"""for internal use, initialize all steps (recursively)"""
if self.secret_sources:
self._secrets = SecretsStore.from_list(self.secret_sources)
if self.error_stream:
self._error_stream_object = get_stream_pusher(self.error_stream)
self.resource_cache = resource_cache or ResourceCache()
context = GraphContext(server=self,
nuclio_context=context,
logger=logger)
context.is_mock = is_mock
context.root = self.graph
context.stream = _StreamContext(self.track_models, self.parameters,
self.function_uri)
context.current_function = self._current_function
context.get_store_resource = self.resource_cache.resource_getter(
self._get_db(), self._secrets)
context.get_table = self.resource_cache.get_table
context.verbose = self.verbose
self.context = context
if self.graph_initializer:
if callable(self.graph_initializer):
handler = self.graph_initializer
else:
handler = get_function(self.graph_initializer, namespace or [])
handler(self)
context.root = self.graph
self.graph.init_object(context, namespace, self.load_mode, reset=True)
return v2_serving_handler
def test_run_with_k8s_secrets(self, db: Session, client: TestClient):
project_secret_name = "dummy_secret_name"
secret_keys = ["secret1", "secret2", "secret3"]
# Need to do some mocking, so code thinks that the secret contains these keys. Otherwise it will not add
# the env. variables to the pod spec.
get_k8s().get_project_secret_name = unittest.mock.Mock(
return_value=project_secret_name)
get_k8s().get_project_secret_keys = unittest.mock.Mock(
return_value=secret_keys)
runtime = self._generate_runtime()
task = self._generate_task()
task.metadata.project = self.project
secret_source = {
"kind": "kubernetes",
"source": secret_keys,
}
task.with_secrets(secret_source["kind"], secret_keys)
# What we expect in this case is that environment variables will be added to the pod which get their
# value from the k8s secret, using the correct keys.
expected_env_from_secrets = {}
for key in secret_keys:
env_variable_name = SecretsStore._k8s_env_variable_name_for_secret(
key)
expected_env_from_secrets[env_variable_name] = {
project_secret_name: key
}
self._execute_run(runtime, runspec=task)
self._assert_pod_creation_config(
expected_secrets=secret_source,
expected_env_from_secrets=expected_env_from_secrets,
)
def _get_secret_or_env(self, key, default=None):
# Project-secrets are mounted as env variables whose name can be retrieved from SecretsStore
return (self._secret(key) or getenv(key)
or getenv(SecretsStore.k8s_env_variable_name_for_secret(key))
or default)