Add integration tests for auth refresh flow, rate limiting, and statistics

- Add refresh token flow tests including token rotation and invalidation
- Add rate limiting integration tests
- Add statistics endpoint integration tests

.env.example

@@ -29,3 +29,10 @@ REGISTRATION_MODE=open
 # Logging Configuration
 LOG_LEVEL=info
 ENVIRONMENT=production
+
+# Backup Configuration
+BACKUP_ENABLED=false
+BACKUP_INTERVAL=24h
+BACKUP_RETENTION_DAYS=7
+BACKUP_PATH=./data/backups
+BACKUP_COMPRESS=true

cmd/api/main.go

@@ -12,6 +12,7 @@ import (
 	"apocapoc-api/internal/application/queries"
 	"apocapoc-api/internal/i18n"
 	"apocapoc-api/internal/infrastructure/auth"
+	"apocapoc-api/internal/infrastructure/backup"
 	"apocapoc-api/internal/infrastructure/config"
 	"apocapoc-api/internal/infrastructure/crypto"
 	"apocapoc-api/internal/infrastructure/email"
@@ -55,6 +56,27 @@ func main() {
 	}
 	defer db.Close()
 
+	backupInterval, err := parseDuration(cfg.BackupInterval)
+	if err != nil {
+		logger.Fatal().Err(err).Msg("Invalid BACKUP_INTERVAL")
+	}
+
+	backupRetentionDays, err := strconv.Atoi(cfg.BackupRetentionDays)
+	if err != nil {
+		logger.Fatal().Err(err).Msg("Invalid BACKUP_RETENTION_DAYS")
+	}
+
+	backupScheduler := backup.NewScheduler(db.Conn(), backup.Config{
+		Enabled:       cfg.BackupEnabled == "true",
+		Interval:      backupInterval,
+		RetentionDays: backupRetentionDays,
+		Path:          cfg.BackupPath,
+		Compress:      cfg.BackupCompress == "true",
+		DatabasePath:  cfg.DBPath,
+	})
+	backupScheduler.Start()
+	defer backupScheduler.Stop()
+
 	jwtExpiryHours, err := parseJWTExpiry(cfg.JWTExpiry)
 	if err != nil {
 		logger.Fatal().Err(err).Msg("Invalid JWT_EXPIRY")

internal/infrastructure/backup/backup.go

@@ -0,0 +1,95 @@
+package backup
+
+import (
+	"compress/gzip"
+	"database/sql"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"time"
+
+	"apocapoc-api/internal/infrastructure/logger"
+)
+
+type Config struct {
+	Enabled       bool
+	Interval      time.Duration
+	RetentionDays int
+	Path          string
+	Compress      bool
+	DatabasePath  string
+}
+
+func CreateBackup(db *sql.DB, config Config) error {
+	if !config.Enabled {
+		return nil
+	}
+
+	if err := os.MkdirAll(config.Path, 0755); err != nil {
+		return fmt.Errorf("failed to create backup directory: %w", err)
+	}
+
+	timestamp := time.Now().Format("20060102_150405")
+	filename := fmt.Sprintf("apocapoc_%s.db", timestamp)
+	backupPath := filepath.Join(config.Path, filename)
+
+	logger.Info().
+		Str("backup_path", backupPath).
+		Msg("Starting database backup")
+
+	if err := backupDatabase(db, backupPath); err != nil {
+		logger.Error().
+			Err(err).
+			Str("backup_path", backupPath).
+			Msg("Backup failed")
+		return fmt.Errorf("backup failed: %w", err)
+	}
+
+	if config.Compress {
+		compressedPath := backupPath + ".gz"
+		if err := compressFile(backupPath, compressedPath); err != nil {
+			logger.Warn().
+				Err(err).
+				Str("backup_path", backupPath).
+				Msg("Compression failed, keeping uncompressed backup")
+		} else {
+			os.Remove(backupPath)
+			backupPath = compressedPath
+		}
+	}
+
+	logger.Info().
+		Str("backup_path", backupPath).
+		Msg("Backup completed successfully")
+
+	return nil
+}
+
+func backupDatabase(db *sql.DB, destPath string) error {
+	_, err := db.Exec(fmt.Sprintf("VACUUM INTO '%s'", destPath))
+	if err != nil {
+		return fmt.Errorf("vacuum into failed: %w", err)
+	}
+	return nil
+}
+
+func compressFile(srcPath, destPath string) error {
+	srcFile, err := os.Open(srcPath)
+	if err != nil {
+		return err
+	}
+	defer srcFile.Close()
+
+	destFile, err := os.Create(destPath)
+	if err != nil {
+		return err
+	}
+	defer destFile.Close()
+
+	gzipWriter := gzip.NewWriter(destFile)
+	defer gzipWriter.Close()
+
+	_, err = io.Copy(gzipWriter, srcFile)
+	return err
+}

internal/infrastructure/backup/backup_test.go

@@ -0,0 +1,169 @@
+package backup
+
+import (
+	"database/sql"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	_ "modernc.org/sqlite"
+)
+
+func TestCreateBackup(t *testing.T) {
+	tempDir := t.TempDir()
+	dbPath := filepath.Join(tempDir, "test.db")
+
+	db, err := sql.Open("sqlite", dbPath)
+	if err != nil {
+		t.Fatalf("Failed to create test database: %v", err)
+	}
+	defer db.Close()
+
+	_, err = db.Exec("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
+	if err != nil {
+		t.Fatalf("Failed to create test table: %v", err)
+	}
+
+	_, err = db.Exec("INSERT INTO test (name) VALUES ('test1'), ('test2')")
+	if err != nil {
+		t.Fatalf("Failed to insert test data: %v", err)
+	}
+
+	backupPath := filepath.Join(tempDir, "backups")
+	config := Config{
+		Enabled:       true,
+		Interval:      24 * time.Hour,
+		RetentionDays: 7,
+		Path:          backupPath,
+		Compress:      false,
+		DatabasePath:  dbPath,
+	}
+
+	err = CreateBackup(db, config)
+	if err != nil {
+		t.Fatalf("CreateBackup failed: %v", err)
+	}
+
+	files, err := os.ReadDir(backupPath)
+	if err != nil {
+		t.Fatalf("Failed to read backup directory: %v", err)
+	}
+
+	if len(files) != 1 {
+		t.Errorf("Expected 1 backup file, got %d", len(files))
+	}
+
+	if len(files) > 0 && filepath.Ext(files[0].Name()) != ".db" {
+		t.Errorf("Expected backup file to have .db extension, got %s", files[0].Name())
+	}
+}
+
+func TestCreateBackupWithCompression(t *testing.T) {
+	tempDir := t.TempDir()
+	dbPath := filepath.Join(tempDir, "test.db")
+
+	db, err := sql.Open("sqlite", dbPath)
+	if err != nil {
+		t.Fatalf("Failed to create test database: %v", err)
+	}
+	defer db.Close()
+
+	_, err = db.Exec("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
+	if err != nil {
+		t.Fatalf("Failed to create test table: %v", err)
+	}
+
+	backupPath := filepath.Join(tempDir, "backups")
+	config := Config{
+		Enabled:       true,
+		Interval:      24 * time.Hour,
+		RetentionDays: 7,
+		Path:          backupPath,
+		Compress:      true,
+		DatabasePath:  dbPath,
+	}
+
+	err = CreateBackup(db, config)
+	if err != nil {
+		t.Fatalf("CreateBackup failed: %v", err)
+	}
+
+	files, err := os.ReadDir(backupPath)
+	if err != nil {
+		t.Fatalf("Failed to read backup directory: %v", err)
+	}
+
+	if len(files) != 1 {
+		t.Errorf("Expected 1 backup file, got %d", len(files))
+	}
+
+	if len(files) > 0 && filepath.Ext(files[0].Name()) != ".gz" {
+		t.Errorf("Expected backup file to have .gz extension, got %s", files[0].Name())
+	}
+}
+
+func TestCreateBackupDisabled(t *testing.T) {
+	tempDir := t.TempDir()
+	dbPath := filepath.Join(tempDir, "test.db")
+
+	db, err := sql.Open("sqlite", dbPath)
+	if err != nil {
+		t.Fatalf("Failed to create test database: %v", err)
+	}
+	defer db.Close()
+
+	backupPath := filepath.Join(tempDir, "backups")
+	config := Config{
+		Enabled:       false,
+		Interval:      24 * time.Hour,
+		RetentionDays: 7,
+		Path:          backupPath,
+		Compress:      false,
+		DatabasePath:  dbPath,
+	}
+
+	err = CreateBackup(db, config)
+	if err != nil {
+		t.Fatalf("CreateBackup failed: %v", err)
+	}
+
+	_, err = os.Stat(backupPath)
+	if !os.IsNotExist(err) {
+		t.Error("Backup directory should not exist when backup is disabled")
+	}
+}
+
+func TestCleanOldBackups(t *testing.T) {
+	tempDir := t.TempDir()
+	backupPath := filepath.Join(tempDir, "backups")
+	os.MkdirAll(backupPath, 0755)
+
+	oldFile := filepath.Join(backupPath, "apocapoc_20200101_120000.db")
+	recentFile := filepath.Join(backupPath, "apocapoc_"+time.Now().Format("20060102_150405")+".db")
+
+	os.WriteFile(oldFile, []byte("old"), 0644)
+	os.WriteFile(recentFile, []byte("recent"), 0644)
+
+	oldTime := time.Now().AddDate(0, 0, -10)
+	os.Chtimes(oldFile, oldTime, oldTime)
+
+	config := Config{
+		Enabled:       true,
+		RetentionDays: 7,
+		Path:          backupPath,
+	}
+
+	err := CleanOldBackups(config)
+	if err != nil {
+		t.Fatalf("CleanOldBackups failed: %v", err)
+	}
+
+	if _, err := os.Stat(oldFile); !os.IsNotExist(err) {
+		t.Error("Old backup file should have been deleted")
+	}
+
+	if _, err := os.Stat(recentFile); err != nil {
+		t.Error("Recent backup file should still exist")
+	}
+}

internal/infrastructure/backup/retention.go

@@ -0,0 +1,80 @@
+package backup
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"apocapoc-api/internal/infrastructure/logger"
+)
+
+func CleanOldBackups(config Config) error {
+	if !config.Enabled {
+		return nil
+	}
+
+	if config.RetentionDays <= 0 {
+		return nil
+	}
+
+	cutoffTime := time.Now().AddDate(0, 0, -config.RetentionDays)
+
+	files, err := os.ReadDir(config.Path)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return err
+	}
+
+	deletedCount := 0
+	for _, file := range files {
+		if file.IsDir() {
+			continue
+		}
+
+		if !strings.HasPrefix(file.Name(), "apocapoc_") {
+			continue
+		}
+
+		if !strings.HasSuffix(file.Name(), ".db") && !strings.HasSuffix(file.Name(), ".db.gz") {
+			continue
+		}
+
+		filePath := filepath.Join(config.Path, file.Name())
+		info, err := os.Stat(filePath)
+		if err != nil {
+			logger.Warn().
+				Err(err).
+				Str("file", filePath).
+				Msg("Failed to stat backup file")
+			continue
+		}
+
+		if info.ModTime().Before(cutoffTime) {
+			if err := os.Remove(filePath); err != nil {
+				logger.Warn().
+					Err(err).
+					Str("file", filePath).
+					Msg("Failed to delete old backup")
+				continue
+			}
+
+			logger.Info().
+				Str("file", file.Name()).
+				Time("mod_time", info.ModTime()).
+				Msg("Deleted old backup")
+			deletedCount++
+		}
+	}
+
+	if deletedCount > 0 {
+		logger.Info().
+			Int("deleted_count", deletedCount).
+			Int("retention_days", config.RetentionDays).
+			Msg("Backup cleanup completed")
+	}
+
+	return nil
+}

internal/infrastructure/backup/scheduler.go

@@ -0,0 +1,77 @@
+package backup
+
+import (
+	"database/sql"
+	"time"
+
+	"apocapoc-api/internal/infrastructure/logger"
+)
+
+type Scheduler struct {
+	db     *sql.DB
+	config Config
+	stopCh chan struct{}
+}
+
+func NewScheduler(db *sql.DB, config Config) *Scheduler {
+	return &Scheduler{
+		db:     db,
+		config: config,
+		stopCh: make(chan struct{}),
+	}
+}
+
+func (s *Scheduler) Start() {
+	if !s.config.Enabled {
+		logger.Info().Msg("Backup scheduler is disabled")
+		return
+	}
+
+	logger.Info().
+		Dur("interval", s.config.Interval).
+		Int("retention_days", s.config.RetentionDays).
+		Str("path", s.config.Path).
+		Bool("compress", s.config.Compress).
+		Msg("Starting backup scheduler")
+
+	go s.run()
+}
+
+func (s *Scheduler) run() {
+	if err := CreateBackup(s.db, s.config); err != nil {
+		logger.Error().Err(err).Msg("Initial backup failed")
+	}
+
+	if err := CleanOldBackups(s.config); err != nil {
+		logger.Error().Err(err).Msg("Initial cleanup failed")
+	}
+
+	ticker := time.NewTicker(s.config.Interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			logger.Debug().Msg("Running scheduled backup")
+
+			if err := CreateBackup(s.db, s.config); err != nil {
+				logger.Error().Err(err).Msg("Scheduled backup failed")
+				continue
+			}
+
+			if err := CleanOldBackups(s.config); err != nil {
+				logger.Error().Err(err).Msg("Backup cleanup failed")
+			}
+
+		case <-s.stopCh:
+			logger.Info().Msg("Backup scheduler stopped")
+			return
+		}
+	}
+}
+
+func (s *Scheduler) Stop() {
+	if s.config.Enabled {
+		close(s.stopCh)
+	}
+}

internal/infrastructure/config/config.go

@@ -25,6 +25,11 @@ type Config struct {
 	RegistrationMode    string
 	LogLevel            string
 	Environment         string
+	BackupEnabled       string
+	BackupInterval      string
+	BackupRetentionDays string
+	BackupPath          string
+	BackupCompress      string
 }
 
 func Load() (*Config, error) {
@@ -48,6 +53,11 @@ func Load() (*Config, error) {
 		RegistrationMode:    getEnvOrDefault("REGISTRATION_MODE", "open"),
 		LogLevel:            getEnvOrDefault("LOG_LEVEL", "info"),
 		Environment:         getEnvOrDefault("ENVIRONMENT", "production"),
+		BackupEnabled:       getEnvOrDefault("BACKUP_ENABLED", "false"),
+		BackupInterval:      getEnvOrDefault("BACKUP_INTERVAL", "24h"),
+		BackupRetentionDays: getEnvOrDefault("BACKUP_RETENTION_DAYS", "7"),
+		BackupPath:          getEnvOrDefault("BACKUP_PATH", "./data/backups"),
+		BackupCompress:      getEnvOrDefault("BACKUP_COMPRESS", "true"),
 	}
 
 	if cfg.DBPath == "" {

internal/infrastructure/http/auth_integration_test.go

@@ -131,3 +131,94 @@ func TestAuthFlow(t *testing.T) {
 		}
 	})
 }
+
+func TestRefreshTokenFlow(t *testing.T) {
+	ts := setupTestServer(t)
+	defer ts.Close()
+
+	t.Run("Complete refresh token flow", func(t *testing.T) {
+		registerBody := RegisterRequest{
+			Email:    "refresh@example.com",
+			Password: "Password123!",
+		}
+		makeRequest(t, *ts.Router, "POST", "/api/v1/auth/register", registerBody, "")
+
+		loginBody := LoginRequest{
+			Email:    "refresh@example.com",
+			Password: "Password123!",
+		}
+		rr := makeRequest(t, *ts.Router, "POST", "/api/v1/auth/login", loginBody, "")
+
+		var loginResp AuthResponse
+		decodeResponse(t, rr, &loginResp)
+
+		if loginResp.RefreshToken == "" {
+			t.Fatal("Expected refresh token in login response")
+		}
+
+		refreshReq := map[string]string{
+			"refresh_token": loginResp.RefreshToken,
+		}
+		rr = makeRequest(t, *ts.Router, "POST", "/api/v1/auth/refresh", refreshReq, "")
+
+		if rr.Code != http.StatusOK {
+			t.Fatalf("Expected status 200, got %d. Body: %s", rr.Code, rr.Body.String())
+		}
+
+		var refreshResp AuthResponse
+		decodeResponse(t, rr, &refreshResp)
+
+		if refreshResp.Token == "" {
+			t.Error("Expected new access token in refresh response")
+		}
+		if refreshResp.RefreshToken == "" {
+			t.Error("Expected new refresh token in refresh response")
+		}
+	})
+
+	t.Run("Refresh with invalid token", func(t *testing.T) {
+		refreshReq := map[string]string{
+			"refresh_token": "invalid-token",
+		}
+		rr := makeRequest(t, *ts.Router, "POST", "/api/v1/auth/refresh", refreshReq, "")
+
+		if rr.Code != http.StatusUnauthorized {
+			t.Errorf("Expected status 401, got %d", rr.Code)
+		}
+	})
+
+	t.Run("Logout invalidates refresh token", func(t *testing.T) {
+		registerBody := RegisterRequest{
+			Email:    "logout@example.com",
+			Password: "Password123!",
+		}
+		makeRequest(t, *ts.Router, "POST", "/api/v1/auth/register", registerBody, "")
+
+		loginBody := LoginRequest{
+			Email:    "logout@example.com",
+			Password: "Password123!",
+		}
+		rr := makeRequest(t, *ts.Router, "POST", "/api/v1/auth/login", loginBody, "")
+
+		var loginResp AuthResponse
+		decodeResponse(t, rr, &loginResp)
+
+		logoutReq := map[string]string{
+			"refresh_token": loginResp.RefreshToken,
+		}
+		rr = makeRequest(t, *ts.Router, "POST", "/api/v1/auth/logout", logoutReq, loginResp.Token)
+
+		if rr.Code != http.StatusOK {
+			t.Fatalf("Expected status 200 for logout, got %d", rr.Code)
+		}
+
+		refreshReq := map[string]string{
+			"refresh_token": loginResp.RefreshToken,
+		}
+		rr = makeRequest(t, *ts.Router, "POST", "/api/v1/auth/refresh", refreshReq, "")
+
+		if rr.Code != http.StatusUnauthorized {
+			t.Errorf("Expected status 401 when using logged out token, got %d", rr.Code)
+		}
+	})
+}

internal/infrastructure/http/ratelimit_integration_test.go

@@ -0,0 +1,69 @@
+package http
+
+import (
+	"net/http"
+	"testing"
+)
+
+func TestGlobalRateLimiting(t *testing.T) {
+	ts := setupTestServer(t)
+	defer ts.Close()
+
+	token := registerAndLogin(t, *ts.Router, "ratelimit@example.com", "Password123!")
+
+	t.Run("Request within rate limit succeeds", func(t *testing.T) {
+		for i := 0; i < 10; i++ {
+			rr := makeRequest(t, *ts.Router, "GET", "/api/v1/habits", nil, token)
+			if rr.Code == http.StatusTooManyRequests {
+				t.Errorf("Request %d hit rate limit unexpectedly", i+1)
+				break
+			}
+		}
+	})
+}
+
+func TestPasswordResetRateLimiting(t *testing.T) {
+	ts := setupTestServer(t)
+	defer ts.Close()
+
+	registerBody := RegisterRequest{
+		Email:    "resetlimit@example.com",
+		Password: "Password123!",
+	}
+	makeRequest(t, *ts.Router, "POST", "/api/v1/auth/register", registerBody, "")
+
+	t.Run("Email-based rate limit for password reset", func(t *testing.T) {
+		resetReq := map[string]string{
+			"email": "resetlimit@example.com",
+		}
+
+		for i := 0; i < 3; i++ {
+			rr := makeRequest(t, *ts.Router, "POST", "/api/v1/auth/forgot-password", resetReq, "")
+			if rr.Code == http.StatusTooManyRequests {
+				t.Fatalf("Request %d hit rate limit too early (limit is 3)", i+1)
+			}
+		}
+
+		rr := makeRequest(t, *ts.Router, "POST", "/api/v1/auth/forgot-password", resetReq, "")
+		if rr.Code != http.StatusTooManyRequests {
+			t.Errorf("Expected status 429 after 4th request, got %d", rr.Code)
+		}
+	})
+
+	t.Run("Different emails have separate rate limits", func(t *testing.T) {
+		registerBody2 := RegisterRequest{
+			Email:    "resetlimit2@example.com",
+			Password: "Password123!",
+		}
+		makeRequest(t, *ts.Router, "POST", "/api/v1/auth/register", registerBody2, "")
+
+		resetReq := map[string]string{
+			"email": "resetlimit2@example.com",
+		}
+
+		rr := makeRequest(t, *ts.Router, "POST", "/api/v1/auth/forgot-password", resetReq, "")
+		if rr.Code == http.StatusTooManyRequests {
+			t.Error("Different email should not be affected by previous email's rate limit")
+		}
+	})
+}

internal/infrastructure/http/stats_integration_test.go

@@ -0,0 +1,159 @@
+package http
+
+import (
+	"net/http"
+	"testing"
+	"time"
+
+	"apocapoc-api/internal/application/queries"
+)
+
+func TestHabitStatsFlow(t *testing.T) {
+	ts := setupTestServer(t)
+	defer ts.Close()
+
+	token := registerAndLogin(t, *ts.Router, "statsuser@example.com", "Password123!")
+
+	habitBody := CreateHabitRequest{
+		Name:      "Meditation",
+		Type:      "BOOLEAN",
+		Frequency: "DAILY",
+	}
+	rr := makeRequest(t, *ts.Router, "POST", "/api/v1/habits", habitBody, token)
+	var habitResp map[string]string
+	decodeResponse(t, rr, &habitResp)
+	habitID := habitResp["id"]
+
+	t.Run("Stats for new habit should be zero", func(t *testing.T) {
+		rr := makeRequest(t, *ts.Router, "GET", "/api/v1/stats/habits/"+habitID, nil, token)
+
+		if rr.Code != http.StatusOK {
+			t.Fatalf("Expected status 200, got %d. Body: %s", rr.Code, rr.Body.String())
+		}
+
+		var stats queries.HabitStatsDTO
+		decodeResponse(t, rr, &stats)
+
+		if stats.TotalCompletions != 0 {
+			t.Errorf("Expected 0 total completions, got %d", stats.TotalCompletions)
+		}
+		if stats.CurrentStreak != 0 {
+			t.Errorf("Expected 0 current streak, got %d", stats.CurrentStreak)
+		}
+		if stats.LongestStreak != 0 {
+			t.Errorf("Expected 0 longest streak, got %d", stats.LongestStreak)
+		}
+	})
+
+	today := time.Now().UTC().Format("2006-01-02")
+
+	t.Run("Stats after marking habit once", func(t *testing.T) {
+		markReq := MarkHabitRequest{
+			ScheduledDate: today,
+		}
+		rr := makeRequest(t, *ts.Router, "POST", "/api/v1/habits/"+habitID+"/mark", markReq, token)
+		if rr.Code != http.StatusOK {
+			t.Fatalf("Failed to mark habit: %d - %s", rr.Code, rr.Body.String())
+		}
+
+		rr = makeRequest(t, *ts.Router, "GET", "/api/v1/stats/habits/"+habitID, nil, token)
+
+		if rr.Code != http.StatusOK {
+			t.Fatalf("Expected status 200, got %d", rr.Code)
+		}
+
+		var stats queries.HabitStatsDTO
+		decodeResponse(t, rr, &stats)
+
+		if stats.TotalCompletions != 1 {
+			t.Errorf("Expected 1 total completion, got %d", stats.TotalCompletions)
+		}
+		if stats.CurrentStreak != 1 {
+			t.Errorf("Expected current streak of 1, got %d", stats.CurrentStreak)
+		}
+		if stats.LongestStreak != 1 {
+			t.Errorf("Expected longest streak of 1, got %d", stats.LongestStreak)
+		}
+	})
+
+	t.Run("Stats after unmarking habit", func(t *testing.T) {
+		rr := makeRequest(t, *ts.Router, "DELETE", "/api/v1/habits/"+habitID+"/entries/"+today, nil, token)
+		if rr.Code != http.StatusOK {
+			t.Fatalf("Failed to unmark habit: %d", rr.Code)
+		}
+
+		rr = makeRequest(t, *ts.Router, "GET", "/api/v1/stats/habits/"+habitID, nil, token)
+
+		var stats queries.HabitStatsDTO
+		decodeResponse(t, rr, &stats)
+
+		if stats.TotalCompletions != 0 {
+			t.Errorf("Expected 0 total completions after unmark, got %d", stats.TotalCompletions)
+		}
+		if stats.CurrentStreak != 0 {
+			t.Errorf("Expected 0 current streak after unmark, got %d", stats.CurrentStreak)
+		}
+	})
+}
+
+func TestHabitUpdateAffectsStats(t *testing.T) {
+	ts := setupTestServer(t)
+	defer ts.Close()
+
+	token := registerAndLogin(t, *ts.Router, "updatestats@example.com", "Password123!")
+
+	habitBody := CreateHabitRequest{
+		Name:      "Running",
+		Type:      "BOOLEAN",
+		Frequency: "DAILY",
+	}
+	rr := makeRequest(t, *ts.Router, "POST", "/api/v1/habits", habitBody, token)
+	var habitResp map[string]string
+	decodeResponse(t, rr, &habitResp)
+	habitID := habitResp["id"]
+
+	today := time.Now().UTC().Format("2006-01-02")
+	markReq := MarkHabitRequest{
+		ScheduledDate: today,
+	}
+	makeRequest(t, *ts.Router, "POST", "/api/v1/habits/"+habitID+"/mark", markReq, token)
+
+	t.Run("Stats remain after updating habit name", func(t *testing.T) {
+		updateReq := UpdateHabitRequest{
+			Name: "Morning Running",
+		}
+		rr := makeRequest(t, *ts.Router, "PUT", "/api/v1/habits/"+habitID, updateReq, token)
+		if rr.Code != http.StatusOK {
+			t.Fatalf("Failed to update habit: %d", rr.Code)
+		}
+
+		rr = makeRequest(t, *ts.Router, "GET", "/api/v1/stats/habits/"+habitID, nil, token)
+
+		var stats queries.HabitStatsDTO
+		decodeResponse(t, rr, &stats)
+
+		if stats.TotalCompletions != 1 {
+			t.Errorf("Expected stats to persist after update, got %d completions", stats.TotalCompletions)
+		}
+	})
+
+	t.Run("Stats remain available after archiving habit", func(t *testing.T) {
+		rr := makeRequest(t, *ts.Router, "DELETE", "/api/v1/habits/"+habitID, nil, token)
+		if rr.Code != http.StatusOK {
+			t.Fatalf("Failed to archive habit: %d", rr.Code)
+		}
+
+		rr = makeRequest(t, *ts.Router, "GET", "/api/v1/stats/habits/"+habitID, nil, token)
+
+		if rr.Code != http.StatusOK {
+			t.Errorf("Expected stats to remain available for archived habit, got %d", rr.Code)
+		}
+
+		var stats queries.HabitStatsDTO
+		decodeResponse(t, rr, &stats)
+
+		if stats.TotalCompletions != 1 {
+			t.Errorf("Expected stats to persist after archiving, got %d completions", stats.TotalCompletions)
+		}
+	})
+}