fix tests

Author: ravenac95

dagster_sqlmesh/asset.py

@@ -17,6 +17,7 @@
 # Define a SQLMesh Asset
 def sqlmesh_assets(
     *,
+    environment: str,
     config: SQLMeshContextConfig,
     name: t.Optional[str] = None,
     dagster_sqlmesh_translator: t.Optional[SQLMeshDagsterTranslator] = None,
@@ -28,7 +29,7 @@ def sqlmesh_assets(
     controller = DagsterSQLMeshController.setup(config)
     if not dagster_sqlmesh_translator:
         dagster_sqlmesh_translator = SQLMeshDagsterTranslator()
-    conversion = controller.to_asset_outs(dagster_sqlmesh_translator)
+    conversion = controller.to_asset_outs(environment, dagster_sqlmesh_translator)
 
     return multi_asset(
         name=name,

dagster_sqlmesh/conftest.py

@@ -126,7 +126,7 @@ def run(
             plan_options["end"] = end
             run_options["end"] = end
 
-        for _context, event in controller.plan_and_run(
+        for event in controller.plan_and_run(
             environment,
             plan_options=plan_options,
             run_options=run_options,

dagster_sqlmesh/controller/dagster.py

@@ -20,9 +20,10 @@ class DagsterSQLMeshController(SQLMeshController):
     """An extension of the sqlmesh controller specifically for dagster use"""
 
     def to_asset_outs(
-        self, translator: SQLMeshDagsterTranslator
+        self, environment: str, translator: SQLMeshDagsterTranslator
     ) -> SQLMeshMultiAssetOptions:
-        with self.context() as context:
+        with self.instance(environment) as instance:
+            context = instance.context
             dag = context.dag
             output = SQLMeshMultiAssetOptions()
             depsMap: t.Dict[str, CoercibleToAssetDep] = {}

dagster_sqlmesh/test_asset.py

@@ -7,6 +7,6 @@
 def test_sqlmesh_context_to_asset_outs(sample_sqlmesh_test_context: SQLMeshTestContext):
     controller = sample_sqlmesh_test_context.create_controller()
     translator = SQLMeshDagsterTranslator()
-    outs = controller.to_asset_outs(translator)
+    outs = controller.to_asset_outs("dev", translator)
     assert len(list(outs.deps)) == 1
     assert len(outs.outs) == 7

sample/dagster_project/definitions.py

@@ -50,7 +50,7 @@ def test_source() -> pl.DataFrame:
     )
 
 
-@sqlmesh_assets(config=sqlmesh_config)
+@sqlmesh_assets(environment="dev", config=sqlmesh_config)
 def sqlmesh_project(context: AssetExecutionContext, sqlmesh: SQLMeshResource):
     yield from sqlmesh.run(context)