Reading settings models according to the environment

root_with_literals.py

import os
import pytest
from pydantic import BaseSettings, HttpUrl, ValidationError
from typing import Union
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


class Context(BaseSettings):
    __root__: Union[LocalContext, ProdContext]  # 2
    
    @classmethod
    def create(cls):
        model = cls.parse_obj({})  # 3
        return model.__root__  # 4


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

    context = Context.create()

    assert context.env == 'local'  # 5
    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'  # 6
    os.environ['EXTERNAL_SERVER_URL'] = 'http://some.echo.serv.er'
    os.environ['EXTERNAL_SERVER_PORT'] = '8080'

    context = Context.create()

    assert context.env == 'prod'  # 6
    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):
        Context.create()