Feature Flag Service - Sticky Bucket + Kill Switch Solution & Tests

python3 test_flags.py -v

"""Feature Flag Service - targeting, sticky bucketing, kill switch.
Sandbox: python3 -c "from flags import FlagStore; ..."
Tests: python3 test_flags.py
stdlib only. Provisional teaching implementation.
"""
from __future__ import annotations

import hashlib
import threading
from dataclasses import dataclass, field
from typing import Any, Optional


@dataclass
class Flag:
    key: str
    enabled: bool = False
    kill_switch: bool = False
    percentage: int = 0  # 0..100 sticky rollout
    segments: set[str] = field(default_factory=set)
    variants: dict[str, Any] = field(default_factory=lambda: {"on": True, "off": False})


@dataclass(frozen=True)
class Evaluation:
    key: str
    enabled: bool
    reason: str
    variant: Any = False
    bucket: Optional[int] = None


class FlagStore:
    """In-memory flag store with sticky percentage bucketing."""

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._flags: dict[str, Flag] = {}

    def upsert(self, flag: Flag) -> None:
        if not flag.key or not flag.key.strip():
            raise ValueError("INVALID_KEY")
        if not 0 <= flag.percentage <= 100:
            raise ValueError("INVALID_PERCENTAGE")
        with self._lock:
            self._flags[flag.key] = Flag(
                key=flag.key,
                enabled=flag.enabled,
                kill_switch=flag.kill_switch,
                percentage=flag.percentage,
                segments=set(flag.segments),
                variants=dict(flag.variants),
            )

    def get(self, key: str) -> Optional[Flag]:
        with self._lock:
            f = self._flags.get(key)
            return None if f is None else Flag(**{**f.__dict__, "segments": set(f.segments), "variants": dict(f.variants)})

    @staticmethod
    def sticky_bucket(flag_key: str, user_id: str) -> int:
        h = hashlib.sha256(f"{flag_key}:{user_id}".encode()).hexdigest()
        return int(h[:8], 16) % 100

    def evaluate(self, key: str, user_id: str, segment: Optional[str] = None) -> Evaluation:
        with self._lock:
            f = self._flags.get(key)
            if f is None:
                return Evaluation(key, False, "FLAG_NOT_FOUND", False, None)
            if f.kill_switch:
                return Evaluation(key, False, "KILL_SWITCH", f.variants.get("off", False), None)
            if not f.enabled:
                return Evaluation(key, False, "DISABLED", f.variants.get("off", False), None)
            if f.segments:
                if not segment or segment not in f.segments:
                    return Evaluation(key, False, "SEGMENT_MISS", f.variants.get("off", False), None)
            bucket = self.sticky_bucket(key, user_id)
            if bucket < f.percentage:
                return Evaluation(key, True, "PERCENTAGE_HIT", f.variants.get("on", True), bucket)
            return Evaluation(key, False, "PERCENTAGE_MISS", f.variants.get("off", False), bucket)

"""Tests for feature flag service."""
from __future__ import annotations

import threading
import unittest

from flags import Evaluation, Flag, FlagStore


class FlagStoreTests(unittest.TestCase):
    def setUp(self) -> None:
        self.s = FlagStore()

    def test_kill_switch_wins(self):
        self.s.upsert(Flag("x", enabled=True, kill_switch=True, percentage=100))
        e = self.s.evaluate("x", "u1")
        self.assertFalse(e.enabled)
        self.assertEqual(e.reason, "KILL_SWITCH")

    def test_sticky_bucket_stable(self):
        a = FlagStore.sticky_bucket("dark-mode", "user-42")
        b = FlagStore.sticky_bucket("dark-mode", "user-42")
        self.assertEqual(a, b)
        self.assertTrue(0 <= a < 100)

    def test_percentage_rollout(self):
        self.s.upsert(Flag("roll", enabled=True, percentage=0))
        self.assertFalse(self.s.evaluate("roll", "u").enabled)
        self.s.upsert(Flag("roll", enabled=True, percentage=100))
        self.assertTrue(self.s.evaluate("roll", "u").enabled)

    def test_segment_gate(self):
        self.s.upsert(Flag("beta", enabled=True, percentage=100, segments={"internal"}))
        self.assertEqual(self.s.evaluate("beta", "u", segment="external").reason, "SEGMENT_MISS")
        self.assertTrue(self.s.evaluate("beta", "u", segment="internal").enabled)

    def test_missing_flag(self):
        e = self.s.evaluate("nope", "u")
        self.assertEqual(e.reason, "FLAG_NOT_FOUND")

    def test_concurrent_upsert_evaluate(self):
        self.s.upsert(Flag("c", enabled=True, percentage=100))
        errors = []

        def worker(i: int):
            try:
                if i % 2 == 0:
                    self.s.upsert(Flag("c", enabled=True, percentage=100))
                else:
                    assert self.s.evaluate("c", f"u{i}").enabled
            except Exception as e:
                errors.append(e)

        threads = [threading.Thread(target=worker, args=(i,)) for i in range(40)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        self.assertEqual(errors, [])


if __name__ == "__main__":
    unittest.main()