From fb3c5c3184028885eed0f447a155ca113f8e3c80 Mon Sep 17 00:00:00 2001
From: Jeevan Opel <jeevan.opel@snowflake.com>
Date: Wed, 4 Dec 2024 19:50:36 -0800
Subject: [PATCH] Fix test

---
 scripts/plugin.py           | 3 ++-
 scripts/proto_codegen.sh    | 2 +-
 tests/test_protoc_plugin.py | 2 ++
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/scripts/plugin.py b/scripts/plugin.py
index c5fd931..15cbf96 100755
--- a/scripts/plugin.py
+++ b/scripts/plugin.py
@@ -31,6 +31,7 @@
 INLINE_OPTIMIZATION = True
 FILE_PATH_PREFIX = "snowflake.telemetry._internal"
 FILE_NAME_SUFFIX = "_marshaler"
+OPENTELEMETRY_PROTO_DIR = os.environ["OPENTELEMETRY_PROTO_DIR"]
 
 # Inline utility functions
 
@@ -293,7 +294,7 @@ def from_descriptor(descriptor: FileDescriptorProto) -> "FileTemplate":
         # In the case of the opentelemetry-proto files, the preamble is the license header
         # Each line of the preamble is prefixed with "//"
         preamble = ""
-        with open(f'/tmp/opentelemetry-proto/{descriptor.name}', "r") as f:
+        with open(f'{OPENTELEMETRY_PROTO_DIR}/{descriptor.name}', "r") as f:
             line = f.readline()
             while line and line.startswith("//"):
                 preamble += line.replace("//", "#", 1)
diff --git a/scripts/proto_codegen.sh b/scripts/proto_codegen.sh
index 9db288f..4c7b08b 100755
--- a/scripts/proto_codegen.sh
+++ b/scripts/proto_codegen.sh
@@ -60,7 +60,7 @@ find opentelemetry/proto/ -regex ".*_marshaler\.py" -exec rm {} +
 
 # generate proto code for all protos
 all_protos=$(find $PROTO_REPO_DIR/ -iname "*.proto")
-python -m grpc_tools.protoc \
+OPENTELEMETRY_PROTO_DIR=$PROTO_REPO_DIR python -m grpc_tools.protoc \
     -I $PROTO_REPO_DIR \
     --plugin=protoc-gen-custom-plugin=$repo_root/scripts/plugin.py \
     --custom-plugin_out=. \
diff --git a/tests/test_protoc_plugin.py b/tests/test_protoc_plugin.py
index b0508d3..c27d923 100644
--- a/tests/test_protoc_plugin.py
+++ b/tests/test_protoc_plugin.py
@@ -60,6 +60,8 @@ def test_protoc_plugin(self):
             proto_file_dir = os.path.dirname(proto_file.name)
             proto_file_name = os.path.basename(proto_file.name)
 
+            os.environ["OPENTELEMETRY_PROTO_DIR"] = proto_file_dir
+
             # Run protoc with custom plugin to generate serialization code for messages
             result = subprocess.run([
                 "python",