Ability to change working directory in unstable module. (#453)

Author: toondaey

CHANGELOG.md

@@ -16,6 +16,7 @@ Changes are grouped as follows
 
 ### Added
 
+* Ability for users to customize the working directory of the extractor using the `--cwd` flag in the CLI.
 * Added a new command-line argument, --log-level (-l) to allow users to override configured log levels for a single run.
 
 ## 7.7.0

cognite/extractorutils/unstable/core/runtime.py

@@ -127,6 +127,13 @@ def _create_argparser(self) -> ArgumentParser:
             action="store_true",
             help="Skip any checks during startup. Useful for debugging, not recommended for production deployments.",
         )
+        argparser.add_argument(
+            "--cwd",
+            nargs=1,
+            type=Path,
+            required=False,
+            help="Set the current working directory for the extractor.",
+        )
 
         return argparser
 
@@ -209,6 +216,17 @@ def _try_get_application_config(
 
         return application_config, current_config_revision
 
+    def _try_change_cwd(self, cwd: Path | None) -> None:
+        if cwd is not None:
+            try:
+                os.chdir(cwd)
+                self.logger.info(f"Changed working directory to {cwd}")
+            except OSError as e:
+                self.logger.critical(f"Could not change working directory to {cwd}: {e}")
+                raise InvalidConfigError(f"Could not change working directory to {cwd}") from e
+
+        self.logger.info(f"Using {os.getcwd()} as working directory")
+
     def _safe_get_application_config(
         self,
         args: Namespace,
@@ -314,6 +332,7 @@ def run(self) -> None:
         self.logger.info(f"Started runtime with PID {os.getpid()}")
 
         try:
+            self._try_change_cwd(args.cwd[0])
             connection_config = load_file(args.connection_config[0], ConnectionConfig)
         except InvalidConfigError as e:
             self.logger.error(str(e))

tests/test_unstable/conftest.py

@@ -18,6 +18,14 @@
 )
 from cognite.extractorutils.unstable.core.base import Extractor, StartupTask, TaskContext
 
+working_dir = os.getcwd()
+
+
+@pytest.fixture(autouse=True)
+def reset_environment() -> Generator[None, None, None]:
+    yield
+    os.chdir(working_dir)
+
 
 @pytest.fixture
 def set_client() -> CogniteClient:

tests/test_unstable/test_runtime.py

@@ -1,3 +1,4 @@
+import os
 import time
 from argparse import Namespace
 from collections.abc import Generator
@@ -119,3 +120,12 @@ def cancel_after_delay() -> None:
 
     assert len(errors["items"]) == 1
     assert "No configuration found for the given integration" in errors["items"][0]["description"]
+
+
+def test_changing_cwd() -> None:
+    runtime = Runtime(TestExtractor)
+    original_cwd = os.getcwd()
+    runtime._try_change_cwd(Path(__file__).parent)
+
+    assert os.getcwd() == str(Path(__file__).parent)
+    assert os.getcwd() != original_cwd