Improve tests for Spec::FromUrl and url_registry for | syntax.

PiperOrigin-RevId: 783835643
Change-Id: I53dcdaacb74998772702a30e1fd5b5243345c519

Author: laramiel

tensorstore/driver/auto/BUILD

@@ -65,6 +65,7 @@ tensorstore_cc_test(
         "//tensorstore:spec",
         "//tensorstore:transaction",
         "//tensorstore/driver/zarr3",
+        "//tensorstore/internal/os:file_util",
         "//tensorstore/internal/testing:json_gtest",
         "//tensorstore/internal/testing:scoped_directory",
         "//tensorstore/kvstore",

tensorstore/driver/url_registry.cc

@@ -16,7 +16,6 @@
 
 #include <stddef.h>
 
-#include <array>
 #include <optional>
 #include <string>
 #include <string_view>
@@ -92,18 +91,21 @@ Result<TransformedDriverSpec> GetTransformedDriverSpecFromUrlImpl(
     return absl::InvalidArgumentError("URL must be non-empty");
   }
   std::string buffer;
-  auto end_of_scheme = url.find(':');
-  if (end_of_scheme == std::string_view::npos) {
-    if constexpr (!RequireColon) {
-      end_of_scheme = url.size();
-      buffer = tensorstore::StrCat(url, ":");
-      url = buffer;
-    } else {
+  std::string_view scheme =
+      *(absl::StrSplit(url, absl::MaxSplits(':', 1)).begin());
+
+  if (scheme.size() == url.size()) {
+    if constexpr (RequireColon) {
       return absl::InvalidArgumentError(tensorstore::StrCat(
           "URL scheme must be specified in ", tensorstore::QuoteString(url)));
+    } else {
+      // Add a colon to the end of `url`; The compound url syntax does not
+      // require a colon after the scheme, however registered drivers do.
+      buffer = tensorstore::StrCat(url, ":");
+      url = buffer;
     }
   }
-  auto scheme = url.substr(0, end_of_scheme);
+
   Handler handler;
   {
     auto& registry = GetUrlSchemeRegistry();
@@ -157,6 +159,10 @@ Result<TransformedDriverSpec> GetTransformedDriverKvStoreAdapterSpecFromUrl(
 
 Result<TransformedDriverSpec> GetTransformedDriverSpecFromUrl(
     std::string_view url) {
+  if (url.empty()) {
+    return absl::InvalidArgumentError("URL must be non-empty");
+  }
+
   std::variant<std::monostate, kvstore::Spec, TransformedDriverSpec> spec;
 
   auto maybe_apply_auto = [&]() -> absl::Status {
@@ -172,8 +178,8 @@ Result<TransformedDriverSpec> GetTransformedDriverSpecFromUrl(
   };
 
   auto apply_component = [&](std::string_view component) -> absl::Status {
-    auto scheme = static_cast<std::array<std::string_view, 1>>(
-        absl::StrSplit(component, ':'))[0];
+    std::string_view scheme =
+        *(absl::StrSplit(component, absl::MaxSplits(':', 1)).begin());
     auto scheme_kind = internal::GetUrlSchemeKind(scheme);
     auto fail = [&]() -> absl::Status {
       if (!scheme_kind) {
@@ -247,7 +253,12 @@ Result<TransformedDriverSpec> GetTransformedDriverSpecFromUrl(
   };
 
   for (std::string_view component : absl::StrSplit(url, '|')) {
-    TENSORSTORE_RETURN_IF_ERROR(apply_component(component));
+    if (auto status = apply_component(component); !status.ok()) {
+      return tensorstore::MaybeAnnotateStatus(
+          std::move(status),
+          tensorstore::StrCat("Parsing spec from url: ",
+                              tensorstore::QuoteString(url)));
+    }
   }
 
   if (std::holds_alternative<std::monostate>(spec)) {
@@ -256,8 +267,11 @@ Result<TransformedDriverSpec> GetTransformedDriverSpecFromUrl(
     return absl::InvalidArgumentError("Non-empty URL must be specified");
   }
 
-  TENSORSTORE_RETURN_IF_ERROR(maybe_apply_auto());
-
+  if (auto status = maybe_apply_auto(); !status.ok()) {
+    return tensorstore::MaybeAnnotateStatus(
+        std::move(status), tensorstore::StrCat("Parsing spec from url: ",
+                                               tensorstore::QuoteString(url)));
+  }
   return std::move(std::get<TransformedDriverSpec>(spec));
 }
 

tensorstore/kvstore/BUILD

@@ -468,3 +468,21 @@ tensorstore_cc_test(
         "@nlohmann_json//:json",
     ],
 )
+
+tensorstore_cc_test(
+    name = "url_registry_test",
+    srcs = ["url_registry_test.cc"],
+    deps = [
+        ":kvstore",
+        "//tensorstore:context",
+        "//tensorstore/internal:intrusive_ptr",
+        "//tensorstore/util:future",
+        "//tensorstore/util:result",
+        "//tensorstore/util:status_testutil",
+        "//tensorstore/util/garbage_collection",
+        "@abseil-cpp//absl/log:absl_log",
+        "@abseil-cpp//absl/status",
+        "@abseil-cpp//absl/strings",
+        "@googletest//:gtest_main",
+    ],
+)

tensorstore/kvstore/kvstore_test.cc

@@ -62,7 +62,7 @@ TEST(KeyValueStoreTest, TensorStoreUrl) {
   EXPECT_THAT(
       kvstore::Spec::FromJson("json:"),
       MatchesStatus(absl::StatusCode::kInvalidArgument,
-                    "\"json\" is a kvstore-based TensorStore URL scheme: "
+                    ".*\"json\" is a kvstore-based TensorStore URL scheme: "
                     "unsupported URL scheme \"json\" in \"json:\""));
 }
 

tensorstore/kvstore/mock_kvstore.h

@@ -15,6 +15,7 @@
 #ifndef TENSORSTORE_KVSTORE_MOCK_KVSTORE_H_
 #define TENSORSTORE_KVSTORE_MOCK_KVSTORE_H_
 
+#include <functional>
 #include <optional>
 #include <utility>
 

tensorstore/kvstore/url_registry.cc

@@ -84,18 +84,21 @@ Result<kvstore::Spec> GetSpecFromUrlImpl(std::string_view url, Arg&&... arg) {
     return absl::InvalidArgumentError("URL must be non-empty");
   }
   std::string buffer;
-  auto end_of_scheme = url.find(':');
-  if (end_of_scheme == std::string_view::npos) {
-    if constexpr (!RequireColon) {
-      end_of_scheme = url.size();
-      buffer = tensorstore::StrCat(url, ":");
-      url = buffer;
-    } else {
+  std::string_view scheme =
+      *(absl::StrSplit(url, absl::MaxSplits(':', 1)).begin());
+
+  if (scheme.size() == url.size()) {
+    if constexpr (RequireColon) {
       return absl::InvalidArgumentError(tensorstore::StrCat(
           "URL scheme must be specified in ", tensorstore::QuoteString(url)));
+    } else {
+      // Add a colon to the end of `url`; The compound url syntax does not
+      // require a colon after the scheme, however registered drivers do.
+      buffer = tensorstore::StrCat(url, ":");
+      url = buffer;
     }
   }
-  auto scheme = url.substr(0, end_of_scheme);
+
   Handler handler;
   {
     auto& registry = internal_kvstore::GetUrlSchemeRegistry();
@@ -118,10 +121,15 @@ Result<kvstore::Spec> GetSpecFromUrlImpl(std::string_view url, Arg&&... arg) {
   }
   TENSORSTORE_ASSIGN_OR_RETURN(
       auto spec, handler(url, std::forward<Arg>(arg)...),
-      tensorstore::MaybeAnnotateStatus(std::move(_),
-                                       tensorstore::StrCat(tensorstore::StrCat(
-                                           "Invalid kvstore URL component ",
-                                           tensorstore::QuoteString(url)))));
+      tensorstore::MaybeAnnotateStatus(
+          std::move(_), tensorstore::StrCat("Invalid kvstore URL component ",
+                                            tensorstore::QuoteString(url))));
+  if (!spec.valid()) {
+    // This should never happen with correct kvstore drivers, but as Spec
+    // permits invalid state, check here and error about it.
+    return absl::InvalidArgumentError(tensorstore::StrCat(
+        "Invalid kvstore URL component ", tensorstore::QuoteString(url)));
+  }
   TENSORSTORE_RETURN_IF_ERROR(
       const_cast<kvstore::DriverSpec&>(*spec.driver).NormalizeSpec(spec.path));
   return spec;
@@ -144,18 +152,24 @@ Result<kvstore::Spec> GetAdapterSpecFromUrl(std::string_view url,
 namespace kvstore {
 
 Result<Spec> Spec::FromUrl(std::string_view url) {
-  auto splitter = absl::StrSplit(url, '|');
-  auto it = splitter.begin();
-  if (it == splitter.end()) {
-    // Note: `absl::StrSplit` returns either zero or one component for
-    // an empty string, depending on whether `url.data() == nullptr`.
+  // Note: `absl::StrSplit` returns either zero or one component for
+  // an empty string, depending on whether `url.data() == nullptr`,
+  // so check for empty explicitly.
+  if (url.empty()) {
     return absl::InvalidArgumentError("URL must be non-empty");
   }
-  TENSORSTORE_ASSIGN_OR_RETURN(auto spec,
-                               internal_kvstore::GetRootSpecFromUrl(*it));
+  auto splitter = absl::StrSplit(url, '|');
+  auto it = splitter.begin();
+  TENSORSTORE_ASSIGN_OR_RETURN(
+      auto spec, internal_kvstore::GetRootSpecFromUrl(*it),
+      tensorstore::MaybeAnnotateStatus(
+          _, tensorstore::StrCat("Parsing spec from url: ", QuoteString(url))));
   while (++it != splitter.end()) {
     TENSORSTORE_ASSIGN_OR_RETURN(
-        spec, internal_kvstore::GetAdapterSpecFromUrl(*it, std::move(spec)));
+        spec, internal_kvstore::GetAdapterSpecFromUrl(*it, std::move(spec)),
+        tensorstore::MaybeAnnotateStatus(
+            _,
+            tensorstore::StrCat("Parsing spec from url: ", QuoteString(url))));
   }
   return spec;
 }