942 new variable user permission team

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:

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"]