A slightly more concise way to read envs with literal types

root_with_literals_improved.py

import os
from typing import Union

import pytest
from pydantic import BaseSettings, HttpUrl, ValidationError, parse_obj_as
from typing_extensions import Literal


class LocalContext(BaseSettings):  # 1
    env: Literal["local"]  # 1
    echo_server_url: str


class ProdContext(BaseSettings):  # 1
    env: Literal["prod"]  # 1
    external_server_url: HttpUrl
    external_server_port: int


Context = Union[LocalContext, ProdContext]  # 2


def test_reads_local_context():
    os.environ.clear()
    os.environ["ENV"] = "local"  # 4
    os.environ["ECHO_SERVER_URL"] = "http://localhost"

    context = parse_obj_as(Context, {})  # 3

    assert context.env == "local"  # 4
    assert context.echo_server_url == "http://localhost"
    assert not hasattr(context, "external_server_url")
    assert not hasattr(context, "external_server_port")


def test_reads_prod_context():
    os.environ.clear()
    os.environ["ENV"] = "prod"  # 5
    os.environ["EXTERNAL_SERVER_URL"] = "http://some.echo.serv.er"
    os.environ["EXTERNAL_SERVER_PORT"] = "8080"

    context = parse_obj_as(Context, {})  # 3

    assert context.env == "prod"  # 5
    assert context.external_server_url == "http://some.echo.serv.er"
    assert not hasattr(context, "echo_server_url")


def test_only_given_literals_are_supported():
    os.environ.clear()
    os.environ["ENV"] = "staging"

    with pytest.raises(ValidationError):  # 6
        parse_obj_as(Context, {})  # 6