From 8bd7e61c8f4736c480138ba314e6de3f3ab2e53e Mon Sep 17 00:00:00 2001
From: Chris Pryer <14341145+cnpryer@users.noreply.github.com>
Date: Sun, 26 Mar 2023 14:27:45 -0400
Subject: [PATCH] Add PYTHONPATH for test command (#492)

---
 src/huak/ops.rs | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/huak/ops.rs b/src/huak/ops.rs
index f9b4e862..fd97eb21 100644
--- a/src/huak/ops.rs
+++ b/src/huak/ops.rs
@@ -403,12 +403,19 @@ pub fn test_project(config: &OperationConfig) -> HuakResult<()> {
     }
     let mut cmd = Command::new(venv.python_path());
     make_venv_command(&mut cmd, &venv)?;
-    cmd.args(["-m", "pytest"]).args(
-        config
-            .trailing_command_parts
-            .as_ref()
-            .unwrap_or(&Vec::new()),
-    );
+    let python_path = if config.workspace_root.join("src").exists() {
+        config.workspace_root.join("src")
+    } else {
+        config.workspace_root.clone()
+    };
+    cmd.args(["-m", "pytest"])
+        .args(
+            config
+                .trailing_command_parts
+                .as_ref()
+                .unwrap_or(&Vec::new()),
+        )
+        .env("PYTHONPATH", python_path);
     terminal.run_command(&mut cmd)
 }