Skip to content
19 changes: 19 additions & 0 deletions cognite/extractorutils/unstable/core/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,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

Expand Down Expand Up @@ -200,6 +207,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,
Expand Down Expand Up @@ -305,6 +323,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))
Expand Down
8 changes: 8 additions & 0 deletions tests/test_unstable/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,14 @@
)
from cognite.extractorutils.unstable.core.base import Extractor

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:
Expand Down
10 changes: 10 additions & 0 deletions tests/test_unstable/test_runtime.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import os
import time
from argparse import Namespace
from collections.abc import Generator
Expand Down Expand Up @@ -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
Loading