updated test and using enums

Signed-off-by: Sreekanth Vadigi <[email protected]>

Author: sreekanth-db

src/databricks/sql/client.py

@@ -67,7 +67,7 @@
 )
 from databricks.sql.telemetry.latency_logger import log_latency
 from databricks.sql.telemetry.models.enums import StatementType
-from databricks.sql.common.http import DatabricksHttpClient, HttpMethod
+from databricks.sql.common.http import DatabricksHttpClient, HttpMethod, UploadType
 
 logger = logging.getLogger(__name__)
 
@@ -773,7 +773,7 @@ def _handle_staging_put_stream(
             headers=headers,
             timeout=300,  # 5 minute timeout
         ) as response:
-            self._validate_staging_http_response(response, "stream upload")
+            self._validate_staging_http_response(response, UploadType.STREAM_UPLOAD.value)
 
     @log_latency(StatementType.SQL)
     def _handle_staging_put(
@@ -793,7 +793,7 @@ def _handle_staging_put(
         with open(local_file, "rb") as fh:
             r = requests.put(url=presigned_url, data=fh, headers=headers)
 
-        self._validate_staging_http_response(r, "file upload")
+        self._validate_staging_http_response(r, UploadType.FILE_UPLOAD.value)
 
     @log_latency(StatementType.SQL)
     def _handle_staging_get(

src/databricks/sql/common/http.py

@@ -25,6 +25,11 @@ class HttpHeader(str, Enum):
     AUTHORIZATION = "Authorization"
 
 
+class UploadType(str, Enum):
+    FILE_UPLOAD = "file_upload"
+    STREAM_UPLOAD = "stream_upload"
+
+
 # Dataclass for OAuthHTTP Response
 @dataclass
 class OAuthResponse:

tests/e2e/common/streaming_put_tests.py

@@ -19,13 +19,14 @@ def test_streaming_put_basic(self, catalog, schema):
 
         # Create test data
         test_data = b"Hello, streaming world! This is test data."
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        filename = f"stream_test_{timestamp}.txt"
+        filename = "streaming_put_test.txt"
         file_path = f"/Volumes/{catalog}/{schema}/e2etests/{filename}"
 
         try:
             with self.connection() as conn:
                 with conn.cursor() as cursor:
+                    self._cleanup_test_file(file_path)
+                    
                     with io.BytesIO(test_data) as stream:
                         cursor.execute(
                             f"PUT '__input_stream__' INTO '{file_path}'",