diff --git a/tests/connectors/test_train.py b/tests/connectors/test_train.py index 6cf8bd01..0807cd29 100644 --- a/tests/connectors/test_train.py +++ b/tests/connectors/test_train.py @@ -5,6 +5,33 @@ import pytest import logging import re +import importlib + + +class FakeTrainResponse: + ok = True + status_code = 200 + + def __init__(self, model_id="test-model-id"): + self.model_id = model_id + self.text = f'{{"model_id":"{model_id}"}}' + + def json(self): + return {"model_id": self.model_id} + + +def mock_train_model_create(monkeypatch, model_id="test-model-id"): + train_module = importlib.import_module("wrangles.train") + calls = [] + + def post(*args, **kwargs): + calls.append((args, kwargs)) + return FakeTrainResponse(model_id) + + monkeypatch.setattr(train_module._auth, "get_access_token", lambda: "test-token") + monkeypatch.setattr(train_module._requests, "post", post) + return calls + class LogCapture(logging.Handler): def __init__(self, *args, **kwargs): @@ -141,26 +168,30 @@ def test_classify_read_four_cols_error(mocker): """ ) -def test_classify_write_logs_new_model_id_integration(caplog): - df = pd.DataFrame({ - 'Example': ['apple', 'banana'], - 'Category': ['fruit', 'fruit'], - 'Notes': ['', ''] - }) - - try: - wrangles.recipe.run( - """ - write: - - train.classify: - name: Test Classify Model - """, - dataframe=df - ) - - assert any(record.message for record in caplog.records if record.levelname == "INFO" and "New classify model created" in record.message) - finally: - _delete_model_from_log(caplog, "New classify model created") +def test_classify_write_logs_new_model_id_integration(caplog, monkeypatch): + calls = mock_train_model_create(monkeypatch, "classify-test-model") + df = pd.DataFrame({ + 'Example': ['apple', 'banana'], + 'Category': ['fruit', 'fruit'], + 'Notes': ['', ''] + }) + + wrangles.recipe.run( + """ + write: + - train.classify: + name: Test Classify Model + """, + dataframe=df + ) + + assert any(record.message for record in caplog.records if record.levelname == "INFO" and "New classify model created" in record.message) + assert calls[0][1]["params"] == {"type": "classify", "name": "Test Classify Model"} + assert calls[0][1]["json"] == [ + ["Example", "Category", "Notes"], + ["apple", "fruit", ""], + ["banana", "fruit", ""], + ] class TestTrainExtract: """ @@ -744,10 +775,11 @@ def test_update_model(self): df = wrangles.recipe.run(recipe, dataframe=df) assert df.iloc[0]['Key'] == 'Rachel' and df.iloc[0]['Value'] == 'Updated Rachel' - def test_upsert_new_model_recipe(self): + def test_upsert_new_model_recipe(self, monkeypatch): """ - Test upsert creates new model when model_id doesn't exist + Test upsert creates a new model request without persisting a real test model. """ + calls = mock_train_model_create(monkeypatch, "lookup-upsert-test-model") df = pd.DataFrame({ 'Key': ['Rachel', 'NewCharacter'], 'Value': ['Updated Rachel', 'New Movie'] @@ -762,19 +794,18 @@ def test_upsert_new_model_recipe(self): variant: key """ - model_id = None - try: - result = wrangles.recipe.run(recipe, dataframe=df) - assert len(result) == 2 - assert 'NewCharacter' in result['Key'].tolist() - assert result['Value'].tolist() == ['Updated Rachel', 'New Movie'] - models = wrangles.data.user.models('lookup') - model = next((m for m in models if m['name'] == model_name), None) - assert model is not None - model_id = model['id'] - finally: - if model_id: - wrangles.train.delete(model_id) + result = wrangles.recipe.run(recipe, dataframe=df) + assert len(result) == 2 + assert 'NewCharacter' in result['Key'].tolist() + assert result['Value'].tolist() == ['Updated Rachel', 'New Movie'] + assert calls[0][1]["params"] == {"type": "lookup", "name": model_name, "variant": "key"} + assert calls[0][1]["json"] == { + "Columns": ["Key", "Value"], + "Data": [ + ["Rachel", "Updated Rachel"], + ["NewCharacter", "New Movie"], + ], + } def test_action_parameter_upsert(self): """ @@ -1421,33 +1452,36 @@ def test_missing_columns_error_message(self, mocker): wrangles.recipe.run(recipe, dataframe=df) -def test_lookup_write_logs_new_model_id(caplog): - """ - Integration test for lookup model creation logging - """ - df = pd.DataFrame({ - 'Key': ['apple', 'banana'], - 'Value': ['fruit', 'fruit'] - }) - - try: - wrangles.recipe.run( - """ - write: - - train.lookup: - name: Test Lookup Model Integration - variant: key - """, - dataframe=df - ) - - # Check that model_id was logged - assert any( - record.message for record in caplog.records - if record.levelname == "INFO" and "New lookup model created" in record.message - ) - finally: - _delete_model_from_log(caplog, "New lookup model created") +def test_lookup_write_logs_new_model_id(caplog, monkeypatch): + """ + Test lookup model creation logging without creating a real test model. + """ + calls = mock_train_model_create(monkeypatch, "lookup-test-model") + df = pd.DataFrame({ + 'Key': ['apple', 'banana'], + 'Value': ['fruit', 'fruit'] + }) + + wrangles.recipe.run( + """ + write: + - train.lookup: + name: Test Lookup Model Integration + variant: key + """, + dataframe=df + ) + + # Check that model_id was logged + assert any( + record.message for record in caplog.records + if record.levelname == "INFO" and "New lookup model created" in record.message + ) + assert calls[0][1]["params"] == {"type": "lookup", "name": "Test Lookup Model Integration", "variant": "key"} + assert calls[0][1]["json"] == { + "Columns": ["Key", "Value"], + "Data": [["apple", "fruit"], ["banana", "fruit"]], + } # @@ -1531,33 +1565,37 @@ def test_standardize_error(): }) ) -def test_standardize_write_logs_new_model_id(caplog): - """ - Integration test for standardize model creation logging - """ - df = pd.DataFrame({ - 'Find': ['ASAP', 'ETA'], - 'Replace': ['As Soon As Possible', 'Estimated Time of Arrival'], - 'Notes': ['', ''] - }) - - try: - wrangles.recipe.run( - """ - write: - - train.standardize: - name: Test Standardize Model Integration - """, - dataframe=df - ) - - # Check that model creation was logged - assert any( - record.message for record in caplog.records - if record.levelname == "INFO" and "Creating new standardize model" in record.message - ) - finally: - _delete_model_from_log(caplog, "New standardize model created") +def test_standardize_write_logs_new_model_id(caplog, monkeypatch): + """ + Test standardize model creation logging without creating a real test model. + """ + calls = mock_train_model_create(monkeypatch, "standardize-test-model") + df = pd.DataFrame({ + 'Find': ['ASAP', 'ETA'], + 'Replace': ['As Soon As Possible', 'Estimated Time of Arrival'], + 'Notes': ['', ''] + }) + + wrangles.recipe.run( + """ + write: + - train.standardize: + name: Test Standardize Model Integration + """, + dataframe=df + ) + + # Check that model creation was logged + assert any( + record.message for record in caplog.records + if record.levelname == "INFO" and "Creating new standardize model" in record.message + ) + assert calls[0][1]["params"] == {"type": "standardize", "name": "Test Standardize Model Integration"} + assert calls[0][1]["json"] == [ + ["Find", "Replace", "Notes"], + ["ASAP", "As Soon As Possible", ""], + ["ETA", "Estimated Time of Arrival", ""], + ] class TestTrainMetaData: diff --git a/tests/recipes/test_recipes.py b/tests/recipes/test_recipes.py index b2e09ef5..38097262 100644 --- a/tests/recipes/test_recipes.py +++ b/tests/recipes/test_recipes.py @@ -127,11 +127,36 @@ def test_recipe_by_version_tag(): ) -def test_recipe_by_production_version(): +def test_recipe_by_production_version(mocker): """ Test running a recipe using a model ID and production version """ + mocker.patch( + "wrangles.data.model", + return_value={ + "purpose": "recipe", + "production_version_id": "production-version-id" + } + ) + model_content = mocker.patch( + "wrangles.data.model_content", + return_value={ + "recipe": """ + read: + - test: + rows: 20 + values: + header: value1 + """ + } + ) + df = wrangles.recipe.run("a6bac9e7-2388-4347") + + model_content.assert_called_once_with( + "a6bac9e7-2388-4347", + "production-version-id" + ) assert ( len(df) == 20 and list(df.columns) == ["header"] @@ -184,21 +209,62 @@ def test_recipe_by_production_semantic_version_falls_back_to_latest( assert "No production version exists, defaulting to latest version" in caplog.text -def test_recipe_by_version_latest(): +def test_recipe_by_version_latest(mocker): """ Test running a recipe using a model ID and latest version """ + mocker.patch( + "wrangles.data.model", + return_value={ + "purpose": "recipe", + "production_version_id": "production-version-id" + } + ) + model_content = mocker.patch( + "wrangles.data.model_content", + return_value={ + "recipe": """ + read: + - test: + rows: 10 + values: + header: value1 + """ + } + ) + df = wrangles.recipe.run("a6bac9e7-2388-4347:latest") + + model_content.assert_called_once_with("a6bac9e7-2388-4347", None) assert ( len(df) == 10 and list(df.columns) == ["header"] ) -def test_recipe_by_latest_version(): +def test_recipe_by_latest_version(mocker): """ Test running a recipe using a model ID and latest version """ + mocker.patch( + "wrangles.data.model", + return_value={"purpose": "recipe"} + ) + model_content = mocker.patch( + "wrangles.data.model_content", + return_value={ + "recipe": """ + read: + - test: + rows: 15 + values: + header: value1 + """ + } + ) + df = wrangles.recipe.run("02fc0c63-1294-415b") + + model_content.assert_called_once_with("02fc0c63-1294-415b", None) assert ( len(df) == 15 and list(df.columns) == ["header"] diff --git a/tests/recipes/test_variables.py b/tests/recipes/test_variables.py index 6529a86f..d4cebfcb 100644 --- a/tests/recipes/test_variables.py +++ b/tests/recipes/test_variables.py @@ -365,3 +365,104 @@ def test_variables_variable_overwrite(): variables={'recipe_variables': 'This is a string'} ) assert isinstance(df['vars'][0], dict) + + +def test_applied_permission_group_variable(monkeypatch): + """ + Test that the authenticated user's effective permission group is available as a recipe variable. + """ + token = wrangles.auth._jwt.encode( + {"applied_permission_group": "enterprise"}, + "test-secret", + algorithm="HS256" + ) + monkeypatch.setattr(wrangles.auth, "get_access_token", lambda: token) + + df = wrangles.recipe.run( + """ + read: + - test: + rows: 1 + values: + group: ${applied_permission_group} + """ + ) + + assert df['group'][0] == 'enterprise' + + +def test_applied_permission_group_variable_if(monkeypatch): + """ + Test that applied_permission_group can be used in Python-style if conditions. + """ + monkeypatch.setattr(wrangles.auth, "get_applied_permission_group", lambda: "enterprise") + + df = wrangles.recipe.run( + """ + read: + - test: + rows: 1 + values: + result: kept + wrangles: + - create.column: + output: allowed + value: true + if: applied_permission_group == 'enterprise' + """ + ) + + assert df['allowed'][0] == True + + +def test_applied_permission_group_variable_user_override(monkeypatch): + """ + Test that explicit variables still override the authenticated permission group. + """ + monkeypatch.setattr(wrangles.auth, "get_applied_permission_group", lambda: "enterprise") + + df = wrangles.recipe.run( + """ + read: + - test: + rows: 1 + values: + group: ${applied_permission_group} + """, + variables={"applied_permission_group": "manual"} + ) + + assert df['group'][0] == 'manual' + + +def test_applied_permission_group_variable_from_recipe_metadata(monkeypatch): + """ + Test that recipe metadata permission group is preferred for remote recipes. + """ + monkeypatch.setattr(wrangles.auth, "get_applied_permission_group", lambda: "token-group") + monkeypatch.setattr( + wrangles.recipe._data, + "model", + lambda model_id: { + "purpose": "recipe", + "production_version_id": "v1", + "applied_permission_group": "metadata-group", + } + ) + monkeypatch.setattr( + wrangles.recipe._data, + "model_content", + lambda model_id, version_id=None: { + "recipe": """ + read: + - test: + rows: 1 + values: + group: ${applied_permission_group} + """ + } + ) + + df = wrangles.recipe.run("12345678-1234-1234") + + assert df["group"][0] == "metadata-group" diff --git a/wrangles/auth.py b/wrangles/auth.py index 31c5adf0..ffda8510 100644 --- a/wrangles/auth.py +++ b/wrangles/auth.py @@ -84,3 +84,33 @@ def get_access_token(): _access_token_expiry = _datetime.now() + _timedelta(0, response.json()['expires_in'] - 30) return _access_token + + +def extract_applied_permission_group(source: dict): + """ + Extract the effective permission group from a metadata or token payload. + """ + if not isinstance(source, dict): + return None + + return source.get("applied_permission_group") + + +def get_applied_permission_group(): + """ + Return the authenticated user's effective permission group from the current access token. + + If no user is authenticated or the token does not contain the claim, + return None so recipes can still run without backend credentials. + """ + try: + token = get_access_token() + except Exception: + return None + + try: + claims = _jwt.decode(token, options={"verify_signature": False}) + except Exception: + return None + + return extract_applied_permission_group(claims) diff --git a/wrangles/recipe.py b/wrangles/recipe.py index 77ed87dc..c35f19e4 100644 --- a/wrangles/recipe.py +++ b/wrangles/recipe.py @@ -20,6 +20,7 @@ from . import recipe_wrangles as _recipe_wrangles from . import connectors as _connectors from . import data as _data +from . import auth as _auth from .config import ( reserved_word_replacements as _reserved_word_replacements, where_overwrite_output as _where_overwrite_output, @@ -76,6 +77,11 @@ def _load_recipe( if variables is None: variables = {} + user_variable_keys = set(variables.keys()) + + if "applied_permission_group" not in variables: + variables["applied_permission_group"] = _auth.get_applied_permission_group() + # Accept path-like objects (e.g. pathlib.Path) by converting to str if isinstance(recipe, _os.PathLike): recipe = str(recipe) @@ -122,6 +128,11 @@ def _load_recipe( if metadata.get('message', None) == 'error': raise ValueError('Incorrect model_id.\nmodel_id may be wrong or does not exists') + metadata_applied_permission_group = _auth.extract_applied_permission_group(metadata) + if metadata_applied_permission_group is not None: + if "applied_permission_group" not in user_variable_keys: + variables["applied_permission_group"] = metadata_applied_permission_group + # Using model_id in wrong function purpose = metadata['purpose'] if purpose != 'recipe': @@ -1219,6 +1230,7 @@ def run( """ if variables is None: variables = {} + variables = variables.copy() parent_context = _RECIPE_RUN_CONTEXT.get() run_context = {