Pytest testing key-value store - environment variable (outside)

import os
import json

path = os.environ.get('STORE_DIR', os.path.dirname(os.path.abspath(__file__)))
filename = os.path.join(path, 'store.json')

def set(key, value):
    data = _read_data()
    data[key] = value
    _save_data(data)

def get(key):
    data = _read_data()
    return(data.get(key))

def _save_data(data):
    with open(filename, 'w') as fh:
        json.dump(data, fh, sort_keys=True, indent=4)


def _read_data():
    data = {}
    if os.path.exists(filename):
        with open(filename) as fh:
            data = json.load(fh)
    return data



if __name__ == '__main__':
    import sys
    if len(sys.argv) == 3:
        cmd, key = sys.argv[1:]
        if cmd == 'get':
            print(get(key))
            exit(0)

    if len(sys.argv) == 4:
        cmd, key, value = sys.argv[1:]
        if cmd == 'set':
            set(key, value)
            print('SET')
            exit(0)
    print(f"""Usage:
           {sys.argv[0]} set key value
           {sys.argv[0]} get key
    """)

import os
import store

def test_store(tmpdir):
    os.environ['STORE_DIR'] = str(tmpdir) # str expected, not LocalPath
    print(os.environ['STORE_DIR'])
    store.set('color', 'Blue')
    assert store.get('color') == 'Blue'

    store.set('color', 'Red')
    assert store.get('color') == 'Red'

    store.set('size', '42')
    assert store.get('size') == '42'
    assert store.get('color') == 'Red'