How to Mock a Feature Store Online Lookup in Python

import random
import time


class OnlineFeatureStore:
    def __init__(self):
        self.features = {}

    def put(self, entity_id: str, feature_name: str, value):
        key = (entity_id, feature_name)
        self.features[key] = (value, time.time())

    def get(self, entity_id: str, feature_name: str):
        key = (entity_id, feature_name)
        if key not in self.features:
            raise KeyError(f"Feature {feature_name!r} not found for entity {entity_id!r}")
        value, timestamp = self.features[key]
        return {"value": value, "timestamp": timestamp}

    def batch_get(self, entity_ids, feature_names):
        results = {}
        for entity_id in entity_ids:
            results[entity_id] = {}
            for feature_name in feature_names:
                try:
                    results[entity_id][feature_name] = self.get(entity_id, feature_name)
                except KeyError:
                    results[entity_id][feature_name] = None
        return results


if __name__ == "__main__":
    store = OnlineFeatureStore()

    store.put("user_123", "age", 34)
    store.put("user_123", "city", "Berlin")
    store.put("user_456", "age", 28)
    store.put("user_456", "city", "Paris")

    print("Single lookup:")
    print(store.get("user_123", "city"))

    print("\nBatch lookup:")
    batch_result = store.batch_get(
        entity_ids=["user_123", "user_456"],
        feature_names=["age", "city", "premium"],
    )
    for entity, features in batch_result.items():
        print(f"{entity}:")
        for feature, data in features.items():
            print(f"  {feature}: {data}")