Tests for SPCS Azure (#46)

Add a few tests for SPCS Azure, as introduced by [object_store_ffi#28](RelationalAI/object_store_ffi#28)

---------

Co-authored-by: André Guedes <[email protected]>

Project.toml

@@ -1,6 +1,6 @@
 name = "RustyObjectStore"
 uuid = "1b5eed3d-1f46-4baa-87f3-a4a892b23610"
-version = "0.10.0"
+version = "0.11.0"
 
 [deps]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
@@ -20,7 +20,7 @@ ReTestItems = "1"
 Sockets = "1"
 Test = "1"
 julia = "1.8"
-object_store_ffi_jll = "0.10.0"
+object_store_ffi_jll = "0.11.0"
 
 [extras]
 CloudBase = "85eb1798-d7c4-4918-bb13-c944d38e27ed"

src/mock_server.jl

@@ -1,4 +1,6 @@
-using CloudBase: CloudCredentials, AWSCredentials, AbstractStore, AWS
+using CloudBase: CloudBase, CloudCredentials, AbstractStore
+using CloudBase: AWSCredentials, AWS
+using CloudBase: AzureCredentials, Azure
 using JSON3, HTTP, Sockets, Base64
 using RustyObjectStore: SnowflakeConfig, ClientOptions
 using Base: UUID
@@ -118,6 +120,31 @@ function construct_stage_info(credentials::AWSCredentials, store::AWS.Bucket, pa
     )
 end
 
+function construct_stage_info(credentials::AzureCredentials, store::Azure.Container, encrypted::Bool)
+    m = match(r"(https?://.*?)/(.*)", store.baseurl)
+    @assert !isnothing(m)
+    test_endpoint = m.captures[1]
+    rest = split(HTTP.unescapeuri(m.captures[2]), "/")
+    account = rest[1]
+    container = rest[2]
+
+    Dict(
+        "locationType" => "AZURE",
+        "location" => container * "/",
+        "path" => container * "/",
+        "region" => "westus2",
+        "storageAccount" => account,
+        "isClientSideEncrypted" => encrypted,
+        "ciphers" => encrypted ? "AES_CBC" : nothing,
+        "creds" => Dict(
+            "AZURE_SAS_TOKEN" => "dummy-token",
+        ),
+        "useS3RegionalUrl" => false,
+        "endPoint" => "blob.core.windows.net",
+        "testEndpoint" => test_endpoint,
+    )
+end
+
 function next_id_and_key(gw::SFGatewayMock)
     @lock gw.keys_lock begin
         key_id = gw.next_key_id
@@ -216,6 +243,8 @@ function start(gw::SFGatewayMock)
 
                 stage_info = if isa(gw.credentials, AWSCredentials) && isa(gw.store, AWS.Bucket)
                     construct_stage_info(gw.credentials, gw.store, stage_path(stage), gw.encrypted)
+                elseif isa(gw.credentials, AzureCredentials) && isa(gw.store, Azure.Container)
+                    construct_stage_info(gw.credentials, gw.store, gw.encrypted)
                 else
                     error("unimplemented")
                 end
@@ -251,18 +280,31 @@ function start(gw::SFGatewayMock)
 
                 stage_info = if isa(gw.credentials, AWSCredentials) && isa(gw.store, AWS.Bucket)
                     construct_stage_info(gw.credentials, gw.store, stage_path(stage), gw.encrypted)
+                elseif isa(gw.credentials, AzureCredentials) && isa(gw.store, Azure.Container)
+                    construct_stage_info(gw.credentials, gw.store, gw.encrypted)
                 else
                     error("unimplemented")
                 end
 
                 encryption_material = if gw.encrypted
-                    # fetch key id from s3 meta and return key
-                    response = AWS.head(
-                        stage_info["testEndpoint"] * "/" * stage_info["location"] * path;
-                        service="s3", region="us-east-1", credentials=gw.credentials
-                    )
-                    pos = findfirst(x -> x[1] == "x-amz-meta-x-amz-matdesc", response.headers)
-                    matdesc = JSON3.read(response.headers[pos][2])
+                    # fetch key id from blob meta and return key
+                    headers, metadata_key = if isa(gw.credentials, AWSCredentials)
+                        response = AWS.head(
+                            stage_info["testEndpoint"] * "/" * stage_info["location"] * path;
+                            service="s3", region="us-east-1", credentials=gw.credentials
+                        )
+                        response.headers, "x-amz-meta-x-amz-matdesc"
+                    elseif isa(gw.credentials, AzureCredentials)
+                        response = Azure.head(
+                            stage_info["testEndpoint"] * "/" * stage_info["storageAccount"] * "/" * stage_info["location"] * path;
+                            service="blob", region="westus2", credentials=gw.credentials
+                        )
+                        response.headers, "x-ms-meta-matdesc"
+                    else
+                        error("unknown credentials type: $(typeof(gw.credentials))")
+                    end
+                    pos = findfirst(x -> x[1] == metadata_key, headers)
+                    matdesc = JSON3.read(headers[pos][2])
                     key_id = matdesc["queryId"]
                     key = find_key_by_id(gw, key_id)
                     Dict(

test/basic_unified_tests.jl

@@ -776,7 +776,7 @@ Minio.with(; debug=true, public=true) do conf
 end # Minio.with
 end # @testitem
 
-@testitem "Basic Snowflake Stage usage" setup=[InitializeObjectStore, SnowflakeMock, ReadWriteCases] begin
+@testitem "Basic Snowflake Stage usage: AWS, non-encrypted" setup=[InitializeObjectStore, SnowflakeMock, ReadWriteCases] begin
 using CloudBase.CloudTest: Minio
 using RustyObjectStore: SnowflakeConfig, ClientOptions
 
@@ -793,7 +793,7 @@ Minio.with(; debug=true, public=false) do conf
 end # Minio.with
 end # @testitem
 
-@testitem "Basic Snowflake Stage usage (encrypted)" setup=[InitializeObjectStore, SnowflakeMock, ReadWriteCases] begin
+@testitem "Basic Snowflake Stage usage: AWS, encrypted" setup=[InitializeObjectStore, SnowflakeMock, ReadWriteCases] begin
 using CloudBase.CloudTest: Minio
 using RustyObjectStore: SnowflakeConfig, ClientOptions
 
@@ -809,3 +809,37 @@ Minio.with(; debug=true, public=false) do conf
     end
 end # Minio.with
 end # @testitem
+
+@testitem "Basic Snowflake Stage usage: Azure, non-encrypted" setup=[InitializeObjectStore, SnowflakeMock, ReadWriteCases] begin
+using CloudBase.CloudTest: Azurite
+using RustyObjectStore: SnowflakeConfig, ClientOptions
+
+# For interactive testing, use Azurite.run() instead of Azurite.with()
+# conf, p = Azurite.run(; debug=true, public=false); atexit(() -> kill(p))
+Azurite.with(; debug=true, public=false) do conf
+    credentials, container = conf
+    with(SFGatewayMock(credentials, container, false)) do config::SnowflakeConfig
+        run_read_write_test_cases(config)
+        run_stream_test_cases(config)
+        run_list_test_cases(config)
+        run_sanity_test_cases(config)
+    end
+end # Azurite.with
+end # @testitem
+
+@testitem "Basic Snowflake Stage usage: Azure, encrypted" setup=[InitializeObjectStore, SnowflakeMock, ReadWriteCases] begin
+using CloudBase.CloudTest: Azurite
+using RustyObjectStore: SnowflakeConfig, ClientOptions
+
+# For interactive testing, use Azurite.run() instead of Azurite.with()
+# conf, p = Azurite.run(; debug=true, public=false); atexit(() -> kill(p))
+Azurite.with(; debug=true, public=false) do conf
+    credentials, container = conf
+    with(SFGatewayMock(credentials, container, true)) do config::SnowflakeConfig
+        run_read_write_test_cases(config)
+        run_stream_test_cases(config)
+        run_list_test_cases(config; strict_entry_size=false)
+        run_sanity_test_cases(config)
+    end
+end # Azurite.with
+end # @testitem