update unit tests

airbnb · Feb 28, 2024 · 992a0e3 · 992a0e3
1 parent 307fcba
commit 992a0e3
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 2 deletions.
diff --git a/api/py/test/sample/teams.json b/api/py/test/sample/teams.json
@@ -32,6 +32,26 @@
                 "EXECUTOR_MEMORY": "4G",
                 "PARALLELISM": "16"
             }
+        },
+        "dev": {
+            "backfill" : {
+                "EXECUTOR_CORES": "1",
+                "DRIVER_MEMORY": "15G",
+                "EXECUTOR_MEMORY": "16G",
+                "PARALLELISM": "4000",
+                "MAX_EXECUTORS": "1000"
+            },
+            "upload" : {
+                "EXECUTOR_CORES": "1",
+                "EXECUTOR_MEMORY": "8G",
+                "PARALLELISM": "1000",
+                "MAX_EXECUTORS": "1000"
+            },
+            "streaming" : {
+                "EXECUTOR_CORES": "2",
+                "EXECUTOR_MEMORY": "4G",
+                "PARALLELISM": "16"
+            }
         }
     },
     "sample_team": {
@@ -42,6 +62,12 @@
             "backfill" : {
                 "EXECUTOR_CORES": "4"
             }
+        },
+        "dev": {
+            "backfill" : {
+                "EXECUTOR_CORES": "2",
+                "DRIVER_MEMORY": "30G"
+            }
         }
     },
     "kaggle": {

diff --git a/api/py/test/test_run.py b/api/py/test/test_run.py
@@ -31,7 +31,7 @@
 def parser():
     """ Basic parser for tests relative to the main arguments of run.py """
     parser = argparse.ArgumentParser()
-    args = ['repo', 'conf', 'mode', 'app-name', 'chronon-jar', 'online-jar', 'online-class', 'render-info', 'sub-help']
+    args = ['repo', 'conf', 'mode', 'env', 'app-name', 'chronon-jar', 'online-jar', 'online-class', 'render-info', 'sub-help']
     for arg in args:
         parser.add_argument(f"--{arg}")
     run.set_defaults(parser)
@@ -98,6 +98,7 @@ def test_environment(teams_json, repo, parser, test_conf_location):
         '--mode', 'backfill',
         '--conf', test_conf_location,
         '--repo', repo,
+        '--env', 'production',
         '--online-jar', test_conf_location,
     ]))
     # from team env.
@@ -111,12 +112,28 @@ def test_environment(teams_json, repo, parser, test_conf_location):
     # from additional_args
     assert os.environ['CHRONON_CONFIG_ADDITIONAL_ARGS'] == '--step-days 14'
 
+    # Check dev backfill for a team sets parameters accordingly.
+    reset_env(default_environment)
+    run.set_runtime_env(parser.parse_args(args=[
+        '--mode', 'backfill',
+        '--conf', test_conf_location,
+        '--repo', repo,
+        '--online-jar', test_conf_location,
+    ]))
+    # from team dev env.
+    assert os.environ['EXECUTOR_CORES'] == '2'
+    # from team dev env.
+    assert os.environ['DRIVER_MEMORY'] == '30G'
+    # from default dev env.
+    assert os.environ['EXECUTOR_MEMORY'] == '16G'
+
     # Check conf set environment overrides most.
     reset_env(default_environment)
     run.set_runtime_env(parser.parse_args(args=[
         '--mode', 'backfill',
         '--conf', 'production/joins/sample_team/sample_join.v1',
-        '--repo', repo
+        '--repo', repo,
+        '--env', 'production',
     ]))
     # from conf env.
     assert os.environ['EXECUTOR_MEMORY'] == '9G'