[MLOps] Monkey Patch

도입

AI 서비스 API 레거시 구현체를 리팩토링 해야할 일을 맡았다.
원래는 타 팀에서 핸들링 하던 코드였는데 RNR이 개편되면서 역할이 우리팀으로 이관되었다.
그래서 리팩토링 뿐만 아니라 배포 절차부터 테스트 파이프라인까지 우리 팀으로 완전 내재화 하기 위한 목적의 백로그를 약 1개월 간 진행했다.

본 글에서는 수행한 여러 작업들 중 유닛 테스트 작성 과정에서 적용한 Python Mock DB Client Monkey Patch 방식에 대해 정리하고자 한다.

본문

시작 단계에서 느꼈던 불편함들이 있었다. 기존 API 내부 로직에서는 데이터베이스에 직접 접근하는데,

• 기존 유닛 테스트는 실제 데이터베이스와 connection을 요구한다.
• 기존 유닛 테스트는 현재 데이터베이스 구조와 호환되지 않는다.

필요한 3개의 DB (mongo, elastic, redis) 중 redis는 테스트 환경에서 따로 띄워줘야 하는 불편함도 있었고 데이터베이스가 정적이지 않아 테스트 쿼리 작성에서 약간의 바틀넥이 있었다. 그래서 이후 테스트 파이프라인에서는 별도 connection을 요구하지 않고 상태가 정적인 Mock DB Client로 갈아 끼워버리기로 했다.

Monkey Patch

: 원래 소스코드를 변경하지 않고 런타임 중 코드 기본 동작을 추가, 변경하는 방식

Monkey Patch Example

아래와 같은 DB Handler가 있다고 하자.

import pandas as pd
import pymongo

class MongoHandler:

  def __init__(self):
    self.client = pymongo.MongoClient(...)
    
  def fetch_as_dataframe(self, query: dict):
    """ 
    몽고디비에서 로우 데이터 조회 후 데이터프레임으로 리턴
    """
    raw = self.client.find(query)
    df = pd.DataFrame(raw)
    return df

이런 DB Handler는 DB client를 갖고 있고 니즈에 맞는 메서드를 붙여서 쓴다.
여기서 DB client를 mock client 객체로 할당한다면 몽키 패치 처리한 것이다.

import pandas as pd
import pymongo
import mongomock

# Monkey Patch
pymongo.MongoClient = mongomock.MongoClient

class MongoHandler:
  def __init__(self):
    self.client = pymongo.MongoClient(...)
    
  def get_dataframe(self, query: dict):
    """ 
    몽고디비에서 로우 데이터 조회 후 데이터프레임으로 리턴
    """
    raw = self.client.find(query)
    df = pd.DataFrame(raw)
    return df

Python Mock DB Client

client는 패치했지만 get_dataframe() 메서드 내부를 보면 client는 .find()라는 인터페이스를 유지하고 있다. 그래서

• mock client는 기존 client와 동일한 인터페이스를 보존해야 한다.

나는 이렇게 인터페이스가 보존되는 mock client 3개를 사용했다.

1. pymongo.MongoClient는 mongomock.MongoClient로 패치
2. elasticsearch.Elasticsearch는 elasticmock.FakeElasticSearch로 패치
3. redis.Redis는 fakeredis.FakeRedis로 패치

위 mock client 들은 실제 client와 거의 유사한 인터페이스를 지원하지만, 간혹 호환되지 않는 부분들이 있을 수 있다. 그런 경우 호환되지 않는 메서드만 오버라이드 해서 맞춤 구현해주면 된다.

도식으로 표현하면 아래와 같다.

Pytest Monkey Patch

pytest로 unittest 를 구현한다면, monkey patch를 아래의 방식으로 적용할 수 있다.

import pytest
# 실제 client
import redis
import pymongo
import elasticsearch
# mock client
import fakeredis
import mongomock
import elasticmock

@pytest.fixture(autouse=True)
def client_monkey_patch(monkeypatch):
    # MongoDB 클라이언트 몽키 패치
    monkeypatch.setattr("pymongo.MongoClient", mongomock.MongoClient)
    # Elasticsearch 클라이언트 몽키 패치
    monkeypatch.setattr("elasticsearch.Elasticsearch", elasticmock.FakeElasticSearch)
    # Redis 클라이언트 몽키 패치
    monkeypatch.setattr("redis.Redis", fakeredis.FakeRedis)

monkeypatch라는 pytest 내장 fixture를 입력으로 받아서 setattr 메서드를 통해 패치 할 수 있다.
나는 conftest.py 파일에 패치 함수를 fixture로 구현하고 test 시작 시점에 패치 하도록 했다.
이때, autouse 옵션을 true로 주면 테스트 세션이 끝나고 패치를 원복한다. 따라서, 별도의 tear-down 로직을 구현하지 않아도 된다.