add optional tf pytest marker, option test_dragonnet

tests/conftest.py

@@ -66,3 +66,23 @@ def _generate_data():
         return data
 
     yield _generate_data
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--runtf", action="store_true", default=False, help="run tf tests"
+    )
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "tf: mark test as tf to run")
+
+
+def pytest_collection_modifyitems(config, items):
+    if config.getoption("--runtf"):
+        # --runtf given in cli: do not skip tf tests
+        return
+    skip_tf = pytest.mark.skip(reason="need --runtf option to run")
+    for item in items:
+        if "tf" in item.keywords:
+            item.add_marker(skip_tf)
+

tests/test_dragonnet.py

@@ -1,8 +1,9 @@
 from causalml.inference.tf import DragonNet
 from causalml.dataset.regression import simulate_nuisance_and_easy_treatment
 import shutil
+import pytest
 
-
+@pytest.mark.tf
 def test_save_load_dragonnet():
     y, X, w, tau, b, e = simulate_nuisance_and_easy_treatment(n=1000)