claude-code-proxy/proxy/internal/service/storage_contract_test.go

package service

import (
	"encoding/json"
	"path/filepath"
	"testing"
	"time"

	"github.com/seifghazi/claude-code-monitor/internal/config"
	"github.com/seifghazi/claude-code-monitor/internal/model"
)

type storageFactory struct {
	name string
	new  func(t *testing.T, cfg config.StorageConfig) StorageService
}

func runStorageContractTests(t *testing.T, factory storageFactory) {
	t.Helper()

	t.Run("save and fetch by short id", func(t *testing.T) {
		storage := factory.new(t, config.StorageConfig{
			DBPath: filepath.Join(t.TempDir(), factory.name+".db"),
		})

		req := newContractRequest("fetch-123")
		mustSaveRequest(t, storage, req)

		got := mustGetByShortID(t, storage, "123")
		if got.RequestID != req.RequestID {
			t.Fatalf("expected request id %q, got %q", req.RequestID, got.RequestID)
		}
		if got.Method != req.Method || got.Endpoint != req.Endpoint || got.Model != req.Model {
			t.Fatalf("unexpected fetched request: %#v", got)
		}
	})

	t.Run("update response persists status and usage metadata", func(t *testing.T) {
		storage := factory.new(t, config.StorageConfig{
			DBPath: filepath.Join(t.TempDir(), factory.name+".db"),
		})

		req := newContractRequest("response-123")
		mustSaveRequest(t, storage, req)
		req.Response = newContractResponse()

		if err := storage.UpdateRequestWithResponse(req); err != nil {
			t.Fatalf("UpdateRequestWithResponse() error = %v", err)
		}

		got := mustGetByShortID(t, storage, "123")
		if got.Response == nil || got.Response.StatusCode != 200 {
			t.Fatalf("expected stored response, got %#v", got.Response)
		}
	})

	t.Run("redaction survives round trip", func(t *testing.T) {
		storage := factory.new(t, config.StorageConfig{
			DBPath:              filepath.Join(t.TempDir(), factory.name+".db"),
			CaptureRequestBody:  true,
			CaptureResponseBody: true,
			RedactedFields:      []string{"api_key", "secret"},
		})

		req := newContractRequest("redact-123")
		req.Body = map[string]interface{}{
			"api_key": "abc123",
			"nested": map[string]interface{}{
				"secret": "top-secret",
				"keep":   "ok",
			},
		}
		req.Response = &model.ResponseLog{
			StatusCode:   httpStatusOK,
			Headers:      map[string][]string{"Content-Type": {"application/json"}},
			Body:         json.RawMessage(`{"secret":"response-secret","visible":"yes"}`),
			ResponseTime: 12,
			CompletedAt:  time.Now().UTC().Format(time.RFC3339),
		}

		mustSaveRequest(t, storage, req)
		if err := storage.UpdateRequestWithResponse(req); err != nil {
			t.Fatalf("UpdateRequestWithResponse() error = %v", err)
		}

		got := mustGetByShortID(t, storage, "123")
		body := got.Body.(map[string]interface{})
		if body["api_key"] != redactionPlaceholder {
			t.Fatalf("expected api_key redacted, got %#v", body["api_key"])
		}
		nested := body["nested"].(map[string]interface{})
		if nested["secret"] != redactionPlaceholder || nested["keep"] != "ok" {
			t.Fatalf("unexpected nested redaction result: %#v", nested)
		}
	})

	t.Run("body suppression semantics", func(t *testing.T) {
		storage := factory.new(t, config.StorageConfig{
			DBPath:              filepath.Join(t.TempDir(), factory.name+".db"),
			CaptureRequestBody:  false,
			CaptureResponseBody: false,
		})

		req := newContractRequest("suppress-123")
		req.Body = map[string]interface{}{"message": "do not store me"}
		req.Response = &model.ResponseLog{
			StatusCode:      httpStatusOK,
			Headers:         map[string][]string{"Content-Type": {"application/json"}},
			Body:            json.RawMessage(`{"answer":"do not store me"}`),
			BodyText:        "sensitive text",
			StreamingChunks: []string{"data: chunk-1"},
			ResponseTime:    10,
			CompletedAt:     time.Now().UTC().Format(time.RFC3339),
		}

		mustSaveRequest(t, storage, req)
		if err := storage.UpdateRequestWithResponse(req); err != nil {
			t.Fatalf("UpdateRequestWithResponse() error = %v", err)
		}

		got := mustGetByShortID(t, storage, "123")
		body := got.Body.(map[string]interface{})
		if body["_storage_mode"] != "request_body_disabled" {
			t.Fatalf("expected request body placeholder, got %#v", body)
		}
		if got.Response == nil || len(got.Response.Body) != 0 || got.Response.BodyText != "" || len(got.Response.StreamingChunks) != 0 {
			t.Fatalf("expected suppressed response body fields, got %#v", got.Response)
		}
	})

	t.Run("retention cleanup on write", func(t *testing.T) {
		storage := factory.new(t, config.StorageConfig{
			DBPath:        filepath.Join(t.TempDir(), factory.name+".db"),
			RetentionDays: 1,
		})

		oldReq := newContractRequest("old-123")
		oldReq.Timestamp = time.Now().Add(-48 * time.Hour).UTC().Format(time.RFC3339)
		mustSaveRequest(t, storage, oldReq)

		recentReq := newContractRequest("recent-123")
		mustSaveRequest(t, storage, recentReq)

		got, err := storage.GetAllRequests("all")
		if err != nil {
			t.Fatalf("GetAllRequests() error = %v", err)
		}
		if len(got) != 1 || got[0].RequestID != "recent-123" {
			t.Fatalf("expected only recent request to remain, got %#v", got)
		}
	})

	t.Run("clear requests removes all rows", func(t *testing.T) {
		storage := factory.new(t, config.StorageConfig{
			DBPath: filepath.Join(t.TempDir(), factory.name+".db"),
		})

		mustSaveRequest(t, storage, newContractRequest("clear-123"))
		mustSaveRequest(t, storage, newContractRequest("clear-456"))

		deleted, err := storage.ClearRequests()
		if err != nil {
			t.Fatalf("ClearRequests() error = %v", err)
		}
		if deleted != 2 {
			t.Fatalf("expected 2 deleted rows, got %d", deleted)
		}

		got, err := storage.GetAllRequests("all")
		if err != nil {
			t.Fatalf("GetAllRequests() error = %v", err)
		}
		if len(got) != 0 {
			t.Fatalf("expected no remaining requests, got %d", len(got))
		}
	})
}

func newContractRequest(id string) *model.RequestLog {
	return &model.RequestLog{
		RequestID:   id,
		Timestamp:   time.Now().UTC().Format(time.RFC3339),
		Method:      "POST",
		Endpoint:    "/v1/messages",
		Headers:     map[string][]string{"Content-Type": {"application/json"}},
		Body:        map[string]interface{}{"message": "hello"},
		Model:       "claude-3-5-sonnet",
		UserAgent:   "test",
		ContentType: "application/json",
	}
}

func newContractResponse() *model.ResponseLog {
	return &model.ResponseLog{
		StatusCode:   httpStatusOK,
		Headers:      map[string][]string{"Content-Type": {"application/json"}},
		Body:         json.RawMessage(`{"usage":{"input_tokens":11,"output_tokens":22},"stop_reason":"end_turn"}`),
		ResponseTime: 17,
		CompletedAt:  time.Now().UTC().Format(time.RFC3339),
	}
}

func mustSaveRequest(t *testing.T, storage StorageService, req *model.RequestLog) {
	t.Helper()
	if _, err := storage.SaveRequest(req); err != nil {
		t.Fatalf("SaveRequest() error = %v", err)
	}
}

func mustGetByShortID(t *testing.T, storage StorageService, shortID string) *model.RequestLog {
	t.Helper()
	got, _, err := storage.GetRequestByShortID(shortID)
	if err != nil {
		t.Fatalf("GetRequestByShortID(%q) error = %v", shortID, err)
	}
	return got
}