Export helper as fixture instead of importing from test package (#147)

tests/benchmark/test_benchmark.py

@@ -7,6 +7,7 @@
 import pandas.testing as pdtest
 import pytest
 from gretel_client.projects.jobs import Status
+from gretel_client.projects.models import read_model_config
 
 from gretel_trainer.benchmark import (
     BenchmarkConfig,
@@ -17,13 +18,66 @@
     create_dataset,
     launch,
 )
-from tests.benchmark.mocks import (
-    DoNothingModel,
-    FailsToGenerate,
-    FailsToTrain,
-    SharedDictLstm,
-    TailoredActgan,
-)
+from gretel_trainer.benchmark.core import Dataset
+from gretel_trainer.benchmark.gretel.models import GretelModel
+
+
+class DoNothingModel:
+    def train(self, source: Dataset, **kwargs) -> None:
+        pass
+
+    def generate(self, **kwargs) -> pd.DataFrame:
+        return pd.DataFrame()
+
+
+class FailsToTrain:
+    def train(self, source: Dataset, **kwargs) -> None:
+        raise Exception("failed")
+
+    def generate(self, **kwargs) -> pd.DataFrame:
+        return pd.DataFrame()
+
+
+class FailsToGenerate:
+    def train(self, source: Dataset, **kwargs) -> None:
+        pass
+
+    def generate(self, **kwargs) -> pd.DataFrame:
+        raise Exception("failed")
+
+
+class TailoredActgan(GretelModel):
+    @property
+    def config(self):
+        c = read_model_config("synthetics/tabular-actgan")
+        c["models"][0]["actgan"]["params"]["epochs"] = 100
+        return c
+
+
+class SharedDictLstm(GretelModel):
+    config = {
+        "schema_version": "1.0",
+        "name": "tabular-lstm",
+        "models": [
+            {
+                "synthetics": {
+                    "data_source": "__tmp__",
+                    "params": {
+                        "epochs": "auto",
+                        "vocab_size": "auto",
+                        "learning_rate": "auto",
+                        "batch_size": "auto",
+                        "rnn_units": "auto",
+                    },
+                    "generate": {"num_records": 5000},
+                    "privacy_filters": {
+                        "outliers": "auto",
+                        "similarity": "auto",
+                    },
+                }
+            }
+        ],
+    }
 
 
 def test_run_with_gretel_dataset(working_dir, project, evaluate_report_path, iris):

tests/relational/conftest.py

@@ -2,7 +2,7 @@
 import sqlite3
 import tempfile
 from pathlib import Path
-from typing import Generator
+from typing import Callable, Generator
 from unittest.mock import Mock, patch
 
 import pandas as pd
@@ -32,8 +32,12 @@ def static_suffix(request):
         yield make_suffix
 
 
-# Doesn't work well as a fixture due to the need for an input param
-def get_invented_table_suffix(make_suffix_execution_number: int):
+@pytest.fixture()
+def get_invented_table_suffix() -> Callable[[int], str]:
+    return _get_invented_table_suffix
+
+
+def _get_invented_table_suffix(make_suffix_execution_number: int):
     return f"invented_{str(make_suffix_execution_number)}"
 
 

tests/relational/test_backup.py

@@ -11,7 +11,6 @@
     BackupSyntheticsTrain,
     BackupTransformsTrain,
 )
-from tests.relational.conftest import get_invented_table_suffix
 
 
 def test_backup_relational_data(trips):
@@ -37,7 +36,7 @@ def test_backup_relational_data(trips):
     assert BackupRelationalData.from_relational_data(trips) == expected
 
 
-def test_backup_relational_data_with_json(documents):
+def test_backup_relational_data_with_json(documents, get_invented_table_suffix):
     purchases_root_invented_table = f"purchases_{get_invented_table_suffix(1)}"
     purchases_data_years_invented_table = f"purchases_{get_invented_table_suffix(2)}"