Include/uncomment dagster-azure integration tests in dagster-databric…

…ks tests (dagster-io#2883) This also unpins some previously hardcoded library versions and alters the environment variable used to run S3/Azure tests so that they can be run separately.
kinghuang · Sep 5, 2020 · d53f5b2 · d53f5b2
1 parent 6e40d7b
commit d53f5b2
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 21 deletions.
diff --git a/python_modules/libraries/dagster-databricks/dagster_databricks_tests/test_pyspark.py b/python_modules/libraries/dagster-databricks/dagster_databricks_tests/test_pyspark.py
@@ -2,6 +2,7 @@
 
 import pytest
 from dagster_aws.s3 import s3_plus_default_intermediate_storage_defs, s3_resource
+from dagster_azure.adls2 import adls2_plus_default_storage_defs, adls2_resource
 from dagster_databricks import databricks_pyspark_step_launcher
 from dagster_pyspark import DataFrame, pyspark_resource
 from pyspark.sql import Row
@@ -20,9 +21,6 @@
 from dagster.seven import mock
 from dagster.utils.merger import deep_merge_dicts
 
-# TODO: uncomment when dagster-azure is merged.
-# from dagster_azure.adls2 import adls2_plus_default_storage_defs, adls2_resource
-
 S3_BUCKET = "dagster-databricks-tests"
 ADLS2_STORAGE_ACCOUNT = "dagsterdatabrickstests"
 ADLS2_CONTAINER = "dagster-databricks-tests"
@@ -42,9 +40,9 @@
             },
         },
         "libraries": [
-            # TODO: uncomment when dagster-azure is merged.
-            # {'pypi': {'package': 'azure-storage-file-datalake~=12.0.1'}},
-            {"pypi": {"package": "dagster-aws==0.7.13"}},
+            {"pypi": {"package": "azure-storage-file-datalake~=12.0.1"}},
+            {"pypi": {"package": "dagster-aws"}},
+            {"pypi": {"package": "dagster-azure"}},
             {"pypi": {"package": "databricks-api"}},
             {"pypi": {"package": "pytest"}},
         ],
@@ -82,16 +80,15 @@ def filter_df_solid(_, people):
 
 
 MODE_DEFS = [
-    # TODO: uncomment when dagster-azure is merged
-    # ModeDefinition(
-    #     'prod_adls2',
-    #     resource_defs={
-    #         'pyspark_step_launcher': databricks_pyspark_step_launcher,
-    #         'pyspark': pyspark_resource,
-    #         'adls2': adls2_resource,
-    #     },
-    #     system_storage_defs=adls2_plus_default_storage_defs,
-    # ),
+    ModeDefinition(
+        "prod_adls2",
+        resource_defs={
+            "pyspark_step_launcher": databricks_pyspark_step_launcher,
+            "pyspark": pyspark_resource,
+            "adls2": adls2_resource,
+        },
+        system_storage_defs=adls2_plus_default_storage_defs,
+    ),
     ModeDefinition(
         "prod_s3",
         resource_defs={
@@ -165,13 +162,12 @@ def test_pyspark_databricks(mock_wait, mock_get_step_events, mock_put_file, mock
     assert result.success
     assert mock_wait.call_count == 1
     assert mock_get_step_events.call_count == 1
-    # TODO: uncomment this with correct value when uploaded packages are more stable
     assert mock_put_file.call_count == 3
     assert mock_submit_run.call_count == 1
 
 
 @pytest.mark.skipif(
-    "DATABRICKS_TEST_DO_IT_LIVE" not in os.environ,
+    "DATABRICKS_TEST_DO_IT_LIVE_S3" not in os.environ,
     reason="This test is slow and requires a Databricks cluster; run only upon explicit request",
 )
 def test_do_it_live_databricks_s3():
@@ -196,9 +192,8 @@ def test_do_it_live_databricks_s3():
     assert result.success
 
 
-@pytest.mark.skip(reason="Need dagster-azure to be merged")
 @pytest.mark.skipif(
-    "DATABRICKS_TEST_DO_IT_LIVE" not in os.environ,
+    "DATABRICKS_TEST_DO_IT_LIVE_ADLS2" not in os.environ,
     reason="This test is slow and requires a Databricks cluster; run only upon explicit request",
 )
 def test_do_it_live_databricks_adls2():
@@ -228,7 +223,7 @@ def test_do_it_live_databricks_adls2():
             "storage": {
                 "adls2": {
                     "config": {
-                        "adls2_container": ADLS2_CONTAINER,
+                        "adls2_file_system": ADLS2_CONTAINER,
                         "adls2_prefix": "dagster-databricks-tests",
                     }
                 }

diff --git a/python_modules/libraries/dagster-databricks/tox.ini b/python_modules/libraries/dagster-databricks/tox.ini
@@ -7,6 +7,7 @@ deps =
   -e ../../dagster
   -r ../../dagster/dev-requirements.txt
   -e ../dagster-aws
+  -e ../dagster-azure
   -e ../dagster-spark
   -e ../dagster-pyspark
   -e .