first commit

tests/async/auth_test.go

@@ -0,0 +1,176 @@
+package async
+
+import (
+	"net/http"
+	"strings"
+	"testing"
+)
+
+// Auth tests cover every combination of VK presence at submit and poll time.
+// All tests use chat_completions as a representative endpoint.
+
+// assertPollSuccess fails the test unless the poll returned a success code (200 or 202).
+func assertPollSuccess(t *testing.T, code int, body []byte) {
+	t.Helper()
+	if code != http.StatusOK && code != http.StatusAccepted {
+		t.Fatalf("expected 200/202, got %d: %s", code, body)
+	}
+}
+
+// TestAuth_Submit_InvalidVK_Returns400 verifies that submitting with a VK value
+// unknown to the governance store fails at submit time with 400.
+// Requires BIFROST_VK to be set, which proves VK governance is active on the server.
+func TestAuth_Submit_InvalidVK_Returns400(t *testing.T) {
+	if cfg.VK == "" {
+		t.Skip("BIFROST_VK not set — governance may not be active")
+	}
+	ec := chatCompletionCase()
+	code, _, body := submitCase(t, ec, vkHeaders("sk-bf-nonexistent-key-for-auth-test"))
+	if code != http.StatusBadRequest {
+		t.Errorf("expected 400 for unknown VK on submit, got %d: %s", code, body)
+	}
+}
+
+// TestAuth_VKScoped_SameKey_Succeeds submits with a VK and polls with the same VK.
+func TestAuth_VKScoped_SameKey_Succeeds(t *testing.T) {
+	if cfg.VK == "" {
+		t.Skip("BIFROST_VK not set")
+	}
+	ec := chatCompletionCase()
+	_, submitted, body := submitCase(t, ec, vkHeaders(cfg.VK))
+	if submitted.ID == "" {
+		t.Fatalf("submit returned no job id: %s", body)
+	}
+
+	pollPath := jobPollPath(ec.pollBase, submitted.ID)
+	code, _, body := pollOnce(t, pollPath, vkHeaders(cfg.VK))
+	assertPollSuccess(t, code, body)
+}
+
+// TestAuth_VKScoped_DifferentKey_Returns404 submits with VK1 and polls with VK2.
+// The gateway must return 404 because the VK IDs will not match.
+func TestAuth_VKScoped_DifferentKey_Returns404(t *testing.T) {
+	if cfg.VK == "" || cfg.AltVK == "" {
+		t.Skip("both BIFROST_VK and BIFROST_ALT_VK must be set")
+	}
+	ec := chatCompletionCase()
+	_, submitted, body := submitCase(t, ec, vkHeaders(cfg.VK))
+	if submitted.ID == "" {
+		t.Fatalf("submit returned no job id: %s", body)
+	}
+
+	pollPath := jobPollPath(ec.pollBase, submitted.ID)
+	code, _, _ := pollOnce(t, pollPath, vkHeaders(cfg.AltVK))
+	if code != http.StatusNotFound {
+		t.Errorf("expected 404 when polling with a different VK, got %d", code)
+	}
+}
+
+// TestAuth_VKScoped_MissingKeyOnPoll_Returns404 submits with a VK and polls
+// without one. The job stores a VirtualKeyID so the gateway requires a VK on poll.
+func TestAuth_VKScoped_MissingKeyOnPoll_Returns404(t *testing.T) {
+	if cfg.VK == "" {
+		t.Skip("BIFROST_VK not set")
+	}
+	ec := chatCompletionCase()
+	_, submitted, body := submitCase(t, ec, vkHeaders(cfg.VK))
+	if submitted.ID == "" {
+		t.Fatalf("submit returned no job id: %s", body)
+	}
+
+	pollPath := jobPollPath(ec.pollBase, submitted.ID)
+	code, _, _ := pollOnce(t, pollPath, nil)
+	if code != http.StatusNotFound {
+		t.Errorf("expected 404 when polling a VK-scoped job without a VK, got %d", code)
+	}
+}
+
+// TestAuth_PublicJob_AnonymousPoll_Succeeds submits without a VK (VirtualKeyID = nil)
+// and polls without a VK. The VK check is skipped for public jobs.
+func TestAuth_PublicJob_AnonymousPoll_Succeeds(t *testing.T) {
+	ec := chatCompletionCase()
+	_, submitted, body := submitCase(t, ec, nil)
+	if submitted.ID == "" {
+		t.Fatalf("submit returned no job id: %s", body)
+	}
+
+	pollPath := jobPollPath(ec.pollBase, submitted.ID)
+	code, _, body := pollOnce(t, pollPath, nil)
+	assertPollSuccess(t, code, body)
+}
+
+// TestAuth_PublicJob_VKPoll_Succeeds submits without a VK and polls with one.
+// Per docs: "Jobs created without a virtual key are not virtual-key scoped, so they
+// can be polled by any caller that passes your gateway auth/middleware checks."
+func TestAuth_PublicJob_VKPoll_Succeeds(t *testing.T) {
+	if cfg.VK == "" {
+		t.Skip("BIFROST_VK not set")
+	}
+	ec := chatCompletionCase()
+	_, submitted, body := submitCase(t, ec, nil)
+	if submitted.ID == "" {
+		t.Fatalf("submit returned no job id: %s", body)
+	}
+
+	pollPath := jobPollPath(ec.pollBase, submitted.ID)
+	code, _, body := pollOnce(t, pollPath, vkHeaders(cfg.VK))
+	assertPollSuccess(t, code, body)
+}
+
+// vkPrefixed returns true when vk begins with the governance virtual-key prefix "sk-bf-".
+// Only keys with this prefix are recognised by the Authorization, x-api-key, and
+// x-goog-api-key header paths in ConvertToBifrostContext.
+func vkPrefixed(vk string) bool {
+	return strings.HasPrefix(strings.ToLower(vk), "sk-bf-")
+}
+
+// TestAuth_BearerVK_SameKey_Succeeds submits with "Authorization: Bearer <vk>" and
+// polls with the same header.  Verifies the Bearer token path in ConvertToBifrostContext.
+func TestAuth_BearerVK_SameKey_Succeeds(t *testing.T) {
+	if cfg.VK == "" || !vkPrefixed(cfg.VK) {
+		t.Skip("BIFROST_VK not set or does not start with sk-bf- prefix")
+	}
+	ec := chatCompletionCase()
+	headers := map[string]string{"Authorization": "Bearer " + cfg.VK}
+	_, submitted, body := submitCase(t, ec, headers)
+	if submitted.ID == "" {
+		t.Fatalf("submit returned no job id: %s", body)
+	}
+	pollPath := jobPollPath(ec.pollBase, submitted.ID)
+	code, _, body := pollOnce(t, pollPath, headers)
+	assertPollSuccess(t, code, body)
+}
+
+// TestAuth_ApiKeyVK_SameKey_Succeeds submits with "x-api-key: <vk>" and polls with
+// the same header.  Verifies the x-api-key path in ConvertToBifrostContext.
+func TestAuth_ApiKeyVK_SameKey_Succeeds(t *testing.T) {
+	if cfg.VK == "" || !vkPrefixed(cfg.VK) {
+		t.Skip("BIFROST_VK not set or does not start with sk-bf- prefix")
+	}
+	ec := chatCompletionCase()
+	headers := map[string]string{"x-api-key": cfg.VK}
+	_, submitted, body := submitCase(t, ec, headers)
+	if submitted.ID == "" {
+		t.Fatalf("submit returned no job id: %s", body)
+	}
+	pollPath := jobPollPath(ec.pollBase, submitted.ID)
+	code, _, body := pollOnce(t, pollPath, headers)
+	assertPollSuccess(t, code, body)
+}
+
+// TestAuth_GoogApiKeyVK_SameKey_Succeeds submits with "x-goog-api-key: <vk>" and polls
+// with the same header.  Verifies the x-goog-api-key path in ConvertToBifrostContext.
+func TestAuth_GoogApiKeyVK_SameKey_Succeeds(t *testing.T) {
+	if cfg.VK == "" || !vkPrefixed(cfg.VK) {
+		t.Skip("BIFROST_VK not set or does not start with sk-bf- prefix")
+	}
+	ec := chatCompletionCase()
+	headers := map[string]string{"x-goog-api-key": cfg.VK}
+	_, submitted, body := submitCase(t, ec, headers)
+	if submitted.ID == "" {
+		t.Fatalf("submit returned no job id: %s", body)
+	}
+	pollPath := jobPollPath(ec.pollBase, submitted.ID)
+	code, _, body := pollOnce(t, pollPath, headers)
+	assertPollSuccess(t, code, body)
+}