From 7854eaf3cef55dd23b718f14948a843fed9246b3 Mon Sep 17 00:00:00 2001
From: abhijeet-dhumal <abhijeetdhumal652@gmail.com>
Date: Wed, 17 Jul 2024 13:04:39 +0530
Subject: [PATCH] Added parameters for S3 compatible storage bucket from which
 to download MNIST datasets

---
 support/environment.go | 34 +++++++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/support/environment.go b/support/environment.go
index 84e1e23..287f4c2 100644
--- a/support/environment.go
+++ b/support/environment.go
@@ -51,6 +51,13 @@ const (
 	// URL for PiPI index containing all the required test Python packages
 	pipIndexURL    = "PIP_INDEX_URL"
 	pipTrustedHost = "PIP_TRUSTED_HOST"
+
+	// Storage bucket credentials
+	storageDefaultEndpoint = "AWS_DEFAULT_ENDPOINT"
+	storageAccessKeyId     = "AWS_ACCESS_KEY_ID"
+	storageSecretKey       = "AWS_SECRET_ACCESS_KEY"
+	storageBucketName      = "AWS_STORAGE_BUCKET"
+	storageBucketMnistDir  = "AWS_STORAGE_BUCKET_MNIST_DIR"
 )
 
 type ClusterType string
@@ -118,7 +125,32 @@ func GetClusterHostname(t Test) string {
 }
 
 func GetMnistDatasetURL() string {
-	return lookupEnvOrDefault(mnistDatasetURL, "http://yann.lecun.com/exdb/mnist/")
+	return lookupEnvOrDefault(mnistDatasetURL, "https://ossci-datasets.s3.amazonaws.com/mnist/")
+}
+
+func GetStorageBucketDefaultEndpoint() (string, bool) {
+	storage_endpoint, exists := os.LookupEnv(storageDefaultEndpoint)
+	return storage_endpoint, exists
+}
+
+func GetStorageBucketAccessKeyId() (string, bool) {
+	storage_access_key_id, exists := os.LookupEnv(storageAccessKeyId)
+	return storage_access_key_id, exists
+}
+
+func GetStorageBucketSecretKey() (string, bool) {
+	storage_secret_key, exists := os.LookupEnv(storageSecretKey)
+	return storage_secret_key, exists
+}
+
+func GetStorageBucketName() (string, bool) {
+	storage_bucket_name, exists := os.LookupEnv(storageBucketName)
+	return storage_bucket_name, exists
+}
+
+func GetStorageBucketMnistDir() (string, bool) {
+	storage_bucket_mnist_dir, exists := os.LookupEnv(storageBucketMnistDir)
+	return storage_bucket_mnist_dir, exists
 }
 
 func GetPipIndexURL() string {