Source code for magpie.mock.network

import json
import os
import pickle
import re

import httpx
from loguru import logger

from magpie import config

CACHE_DIR = f'{config.CACHE_DIR}/http'
os.makedirs(CACHE_DIR, exist_ok=True)




class MockResponse:
    """Mock object for an `httpx.Response`.

    It provides the `text` attribute and the `json()` method.
    """
    def __init__(self, text):
        self.text = text



    def json(self):
        return json.loads(self.text)  # TODO: use msgspec






def slugify(url: str):
    """Return a slug for the given url that can be used as a filename."""
    slug = url.lower()
    slug = re.sub(r'[^a-z0-9]+', '-', slug).strip('-')
    slug = re.sub(r'[-]+', '-', slug)
    return slug



def _cache_response(slug, response):
    """Write the response to disk using the given slug as filename."""
    logger.debug(f'in {CACHE_DIR}: caching {slug}')
    with open(f'{CACHE_DIR}/{slug}', 'wb') as response_file:
        pickle.dump(response, response_file)


def _load_response(slug):
    """Load (if existing) a cached response for the given slug.

    Returns:
        cached content, `None` if no response was found
    """
    try:
        return pickle.load(open(f'{CACHE_DIR}/{slug}', 'rb'))
    except OSError:
        pass




def http_get(url: str, headers=None) -> httpx.Response:
    """Mocked [network.http_get()](#magpie.network.http_get) function.
    """
    slug = slugify(url)

    cached = _load_response(slug)
    if cached is not None:
        logger.debug(f'Returning cached response for {url}')
        return cached.raise_for_status()

    logger.warning(f'Mocked network but no cache (yet) for: {url}')
    headers = headers or {}
    response = httpx.get(url, headers=headers, follow_redirects=True)

    _cache_response(slug, response)

    return response.raise_for_status()