Caso real de refactor

Hoy os traigo un post muy práctico, el otro día me crucé en el trabajo con un código sin testear y bastante acoplado al que le pude aplicar un buen refactor y testing.

Contexto

Como os decía, el otro día trabajando me encontré con un código que estaba bastante acoplado, donde estaban mezcladas diferentes responsabilidades y sobre todo no tenía tests.

El código que tenía que refactorizar era el siguiente:

def model(dbt, session):
    dbt.config(
        packages=["snowflake-snowpark-python", "requests"],
        secrets={"credentials": "token"},
        external_access_integrations=["access_integration"],
    )
    token = _snowflake.get_generic_secret_string(`credentials`)
    headers = {"Authorization": f"Bearer {token}"}
    api_call_url = <API_URL>
    valid_data_response = requests.get(api_call_url, headers = headers)
    if valid_data_response.status_code != 200:
        message = f"ERROR on call to API: Error: {valid_data_response.status_code}."
        raise Exception(message)
    datapoints = valid_data_response.json()[`datapoints`]
    data = pd.DataFrame(datapoints)
    if not data.empty:
        data["value"] = data["value"].astype("str")
    else:
        data = pd.DataFrame({
            "id": pd.Series(dtype="str"),
            "year": pd.Series(dtype="int"),
            "value": pd.Series(dtype="str"),
            "metric": pd.Series(dtype="str")
        })
    return session.create_dataframe(data)

Si le echamos un vistazo podemos ver claramente que la función model hace muchas cosas.

Responsabilidades:

Leer el token de un secreto.
Llamar a una API externa usando ese secreto.
Parsear la respuesta de la API.
Crear un DataFrame con los datos obtenidos.
Llamar a una función de session.

Sin duda muchas responsabilidades para una sola función, vamos a ver como mejorar esto.

Divide y vencerás

Debido a que la función model hace muchas cosas, vamos a mover estas responsabilidades a otras entidades independientes.

Estas pueden ser:

Un cliente HTTP para llamar a la API.
Una clase que recupere los datos válidos de la API y genera el DataFrame.
Nuestra función model que se encargara de llamar a las anteriores.

Cliente HTTP

Como siempre, vamos a empezar por los tests.

Aquí debido a ciertas limitaciones de la API esta no puede ser llamada directamente por lo que nuestros tests validaran:

Que nuestro cliente llama a la URL correcta.
Qué dicha llamada se realiza con los headers necesarios.
Que nuestro cliente es capaz de leer la respuesta y obtener los datos válidos.
Que nuestro cliente es capaz de reacción antes un fallo de la API.

class TestHttpApiClient:
    def test_get_valid_datapoints(self) -> None:
        token = "ANY_TOKEN"
        headers = {"Authorization": f"Bearer {token}"}
        with Stub() as requests:
            datapoints = [{"id": "ANY_ID", "year": 2024, "value": 0, "metric": "ANY_METRIC"}]
            content = {"datapoints": datapoints}
            response = Response()
            response.status_code = OK
            response._content = json.dumps(content).encode()
            requests.get(<API_URL>, headers=headers).returns(response)
        client = HttpApiClient(requests)

        datapoints = client.get_valid_datapoints(token)

        assert datapoints == content["datapoints"]

    def test_raise_error_getting_valid_datapoints(self) -> None:
        token = "ANY_TOKEN"
        headers = {"Authorization": f"Bearer {token}"}
        with Stub() as requests:
            response = Response()
            response.status_code = INTERNAL_SERVER_ERROR
            requests.get(<API_URL>, headers=headers).returns(response)
        client = HttpApiClient(requests)

        expect(lambda: client.get_valid_datapoints(token)).to(raise_error)

Con estas premisas y estos tests nuestro cliente sería algo así:

class HttpApiClient:
    API_URL = <API_URL>

    def __init__(self, _requests=requests) -> None:
        self._requests = _requests

    def get_valid_datapoints(self, token: str) -> list:
        headers = {"Authorization": f"Bearer {token}"}

        valid_data_response = self._requests.get(self.API_URL, headers=headers)

        if valid_data_response.status_code != OK:
            message = f"ERROR on call to API: Error: {valid_data_response.status_code}."
            raise HttpApiclientException(message)

        json_response = valid_data_response.json()
        return json_response["datapoints"]

Podemos ver como nuestro cliente recibe requests como dependencia externa, lo que nos va a permitir mockearlo en nuestros tests y cambiar su comportamiento a nuestro antojo.

Generador de DataFrames

Una vez tenemos nuestro cliente HTTP, vamos a crear una clase que se encargue de generar los DataFrames a partir de los datos de dicha API.

Empiezo por los tests:

class TestGetValidDatapoints:
    def test_get_empty_valid_datapoints(self) -> None:
        token = "ANY_TOKEN"
        datapoints = {"datapoints": []}
        with Stub() as http_api_client:
            http_api_client.get_valid_datapoints().returns(datapoints)
        session = Spy()
        get_valid_datapoints = GetValidDatapoints(http_api_client)

        get_valid_datapoints.execute(token, session)

        expect(session.create_dataframe).to(have_been_called)

    def test_get_valid_datapoints(self) -> None:
        token = "ANY_TOKEN"
        datapoints = [{"id": "ANY_ID", "year": 2024, "value": 0, "metric": "ANY_METRIC"}]
        with Stub() as http_api_client:
            http_api_client.get_valid_datapoints().returns(datapoints)
        session = Spy()
        get_valid_datapoints = GetValidDatapoints(http_api_client)

        get_valid_datapoints.execute(token, session)

        expect(session.create_dataframe).to(have_been_called)

    def test_raise_error_getting_valid_datapoints(self) -> None:
        with Stub() as http_api_client:
            http_api_client.get_valid_datapoints().raises(HttpApiclientException)
        session = Stub()
        get_valid_datapoints = GetValidDatapoints(http_api_client)

        expect(lambda: get_valid_datapoints.execute(token, session)).to(raise_error(GetValidDatapointsException))

Aquí podemos ver como nuestros tests validan que nuestra clase GetValidDatapoints es capaz de generar un DataFrame tanto si la API retorna datos como si no.

De igual forma necesitamos capturar posibles errores en la llamada a la API.

Por tanto, nuestra clase sería algo así:

class GetValidDatapoints:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def execute(self, token: str, session: Session) -> pd.DataFrame:
        try:
            datapoints = self.client.get_valid_datapoints(token)
            data = pd.DataFrame(datapoints)

            if data.empty:
                return session.create_dataframe(self._empty_data)

            data["value"] = data["value"].astype("str")
            return session.create_dataframe(data)
        except HttpApiclientException as ex:
            raise GetValidDatapointsException from ex

    @property
    def _empty_data(self) -> pd.DataFrame:
        return pd.DataFrame(
            {
                "id": pd.Series(dtype="str"),
                "year": pd.Series(dtype="int"),
                "value": pd.Series(dtype="str"),
                "metric": pd.Series(dtype="str"),
            }
        )

Podemos ver como nuestra clase recibe el puerto/interface ApiClient como dependencia externa, lo que nos va a permitir mockearlo en nuestros tests y cambiar su comportamiento a nuestro antojo.

Función `model`

Y por último, el modelo:

def model(dbt, session) -> pd.DataFrame:
    packages=["snowflake-snowpark-python", "requests"]
    secrets={"credentials": "token"}
    integration=["access_integration"]
    dbt.config(packages=packages, secrets=secrets, external_access_integrations=integration)

    client = HttpApiClient()
    get_valid_datapoints = GetValidDatapoints(client)

    token = _snowflake.get_generic_secret_string("credentials")
    return get_valid_datapoints.execute(token, session)

Debido a limitación con DBT nuestra función model no la podemos testear, ya que no podemos cambiar su firma para inyectar nuestras clases, pero bueno como el resto está todo testeado podemos vivir con ello.

Espero que os haya gustado y os haya servido de ayuda, cualquier duda no dudéis en preguntar.

¡Un saludo!