claude-code-proxy/proxy/internal/service/storage_sqlite_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"
)

func TestSQLiteStorageServiceRedactsRequestAndResponseBodies(t *testing.T) {
	storage := newTestSQLiteStorage(t, config.StorageConfig{
		DBPath:              filepath.Join(t.TempDir(), "requests.db"),
		CaptureRequestBody:  true,
		CaptureResponseBody: true,
		RedactedFields:      []string{"api_key", "secret"},
	})

	request := &model.RequestLog{
		RequestID: "redact-123",
		Timestamp: time.Now().UTC().Format(time.RFC3339),
		Method:    "POST",
		Endpoint:  "/v1/messages",
		Headers:   map[string][]string{"Content-Type": {"application/json"}},
		Body: map[string]interface{}{
			"api_key": "abc123",
			"nested": map[string]interface{}{
				"secret":  "top-secret",
				"visible": "ok",
			},
		},
		Model:       "claude-3-5-sonnet",
		UserAgent:   "test",
		ContentType: "application/json",
	}

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

	request.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),
	}

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

	got, _, err := storage.GetRequestByShortID("123")
	if err != nil {
		t.Fatalf("GetRequestByShortID() error = %v", err)
	}

	body, ok := got.Body.(map[string]interface{})
	if !ok {
		t.Fatalf("expected request body to be a map, got %T", got.Body)
	}
	if body["api_key"] != redactionPlaceholder {
		t.Fatalf("expected api_key to be redacted, got %#v", body["api_key"])
	}

	nested, ok := body["nested"].(map[string]interface{})
	if !ok {
		t.Fatalf("expected nested body to be a map, got %T", body["nested"])
	}
	if nested["secret"] != redactionPlaceholder {
		t.Fatalf("expected nested secret to be redacted, got %#v", nested["secret"])
	}
	if nested["visible"] != "ok" {
		t.Fatalf("expected visible field to remain, got %#v", nested["visible"])
	}

	if got.Response == nil || len(got.Response.Body) == 0 {
		t.Fatal("expected response body to be stored")
	}

	var responseBody map[string]interface{}
	if err := json.Unmarshal(got.Response.Body, &responseBody); err != nil {
		t.Fatalf("response body unmarshal failed: %v", err)
	}
	if responseBody["secret"] != redactionPlaceholder {
		t.Fatalf("expected response secret to be redacted, got %#v", responseBody["secret"])
	}
	if responseBody["visible"] != "yes" {
		t.Fatalf("expected response visible field to remain, got %#v", responseBody["visible"])
	}
}

func TestSQLiteStorageServiceHonorsMetadataOnlyMode(t *testing.T) {
	storage := newTestSQLiteStorage(t, config.StorageConfig{
		DBPath:              filepath.Join(t.TempDir(), "requests.db"),
		CaptureRequestBody:  true,
		CaptureResponseBody: true,
		MetadataOnly:        true,
	})

	request := &model.RequestLog{
		RequestID: "metadata-123",
		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": "keep me out of storage",
		},
		Model:       "claude-3-5-sonnet",
		UserAgent:   "test",
		ContentType: "application/json",
		Response: &model.ResponseLog{
			StatusCode: httpStatusOK,
			Headers:    map[string][]string{"Content-Type": {"application/json"}},
			Body:       json.RawMessage(`{"answer":"secret"}`),
			StreamingChunks: []string{
				"data: chunk-1",
			},
			ResponseTime: 10,
			CompletedAt:  time.Now().UTC().Format(time.RFC3339),
		},
	}

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

	got, _, err := storage.GetRequestByShortID("123")
	if err != nil {
		t.Fatalf("GetRequestByShortID() error = %v", err)
	}

	body, ok := got.Body.(map[string]interface{})
	if !ok {
		t.Fatalf("expected metadata-only body placeholder map, got %T", got.Body)
	}
	if body["_storage_mode"] != "metadata_only" {
		t.Fatalf("expected metadata-only placeholder, got %#v", body["_storage_mode"])
	}

	if got.Response == nil {
		t.Fatal("expected response log to exist")
	}
	if len(got.Response.Body) != 0 {
		t.Fatalf("expected response body to be removed, got %s", string(got.Response.Body))
	}
	if got.Response.BodyText != "" {
		t.Fatalf("expected response body text to be removed, got %q", got.Response.BodyText)
	}
	if len(got.Response.StreamingChunks) != 0 {
		t.Fatalf("expected streaming chunks to be removed, got %d", len(got.Response.StreamingChunks))
	}
	if got.Response.StatusCode != httpStatusOK {
		t.Fatalf("expected response status to remain, got %d", got.Response.StatusCode)
	}
}

func TestSQLiteStorageServiceHonorsBodyCaptureToggles(t *testing.T) {
	storage := newTestSQLiteStorage(t, config.StorageConfig{
		DBPath:              filepath.Join(t.TempDir(), "requests.db"),
		CaptureRequestBody:  false,
		CaptureResponseBody: false,
		MetadataOnly:        false,
	})

	request := &model.RequestLog{
		RequestID: "toggle-123",
		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": "do not store me",
		},
		Model:       "claude-3-5-sonnet",
		UserAgent:   "test",
		ContentType: "application/json",
		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),
		},
	}

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

	got, _, err := storage.GetRequestByShortID("123")
	if err != nil {
		t.Fatalf("GetRequestByShortID() error = %v", err)
	}

	body, ok := got.Body.(map[string]interface{})
	if !ok {
		t.Fatalf("expected body placeholder map, got %T", got.Body)
	}
	if body["_storage_mode"] != "request_body_disabled" {
		t.Fatalf("expected request body disabled placeholder, got %#v", body["_storage_mode"])
	}

	if got.Response == nil {
		t.Fatal("expected response log to exist")
	}
	if len(got.Response.Body) != 0 {
		t.Fatalf("expected response body to be omitted, got %s", string(got.Response.Body))
	}
	if got.Response.BodyText != "" {
		t.Fatalf("expected response body text to be omitted, got %q", got.Response.BodyText)
	}
	if len(got.Response.StreamingChunks) != 0 {
		t.Fatalf("expected streaming chunks to be omitted, got %d", len(got.Response.StreamingChunks))
	}
}

func TestSQLiteStorageServiceDeletesExpiredRequestsOnWrite(t *testing.T) {
	storage := newTestSQLiteStorage(t, config.StorageConfig{
		DBPath:         filepath.Join(t.TempDir(), "requests.db"),
		RetentionDays:  1,
		RedactedFields: []string{},
	})

	oldRequest := &model.RequestLog{
		RequestID:   "old-123",
		Timestamp:   time.Now().Add(-48 * time.Hour).UTC().Format(time.RFC3339),
		Method:      "POST",
		Endpoint:    "/v1/messages",
		Headers:     map[string][]string{"Content-Type": {"application/json"}},
		Body:        map[string]interface{}{"message": "old"},
		Model:       "claude-3-5-sonnet",
		UserAgent:   "test",
		ContentType: "application/json",
	}
	if _, err := storage.SaveRequest(oldRequest); err != nil {
		t.Fatalf("SaveRequest(old) error = %v", err)
	}

	recentRequest := &model.RequestLog{
		RequestID:   "recent-123",
		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": "recent"},
		Model:       "claude-3-5-sonnet",
		UserAgent:   "test",
		ContentType: "application/json",
	}
	if _, err := storage.SaveRequest(recentRequest); err != nil {
		t.Fatalf("SaveRequest(recent) error = %v", err)
	}

	got, err := storage.GetAllRequests("all")
	if err != nil {
		t.Fatalf("GetAllRequests() error = %v", err)
	}

	if len(got) != 1 {
		t.Fatalf("expected 1 request after retention cleanup, got %d", len(got))
	}
	if got[0].RequestID != "recent-123" {
		t.Fatalf("expected recent request to remain, got %s", got[0].RequestID)
	}
}

func newTestSQLiteStorage(t *testing.T, cfg config.StorageConfig) *sqliteStorageService {
	t.Helper()

	storage, err := NewSQLiteStorageService(&cfg)
	if err != nil {
		t.Fatalf("NewSQLiteStorageService() error = %v", err)
	}

	sqliteStorage, ok := storage.(*sqliteStorageService)
	if !ok {
		t.Fatalf("unexpected storage type %T", storage)
	}
	t.Cleanup(func() {
		if err := sqliteStorage.Close(); err != nil {
			t.Errorf("Close() error = %v", err)
		}
	})

	return sqliteStorage
}

const httpStatusOK = 200