mirror of
https://github.com/kevin1024/vcrpy.git
synced 2025-12-08 16:53:23 +00:00
135 lines
3.9 KiB
Python
135 lines
3.9 KiB
Python
"""
|
|
Migration script for old 'yaml' and 'json' cassettes
|
|
|
|
.. warning:: Backup your cassettes files before migration.
|
|
|
|
It merges and deletes the request obsolete keys (protocol, host, port, path)
|
|
into new 'uri' key.
|
|
Usage::
|
|
|
|
python -m vcr.migration PATH
|
|
|
|
The PATH can be path to the directory with cassettes or cassette itself
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
import yaml
|
|
|
|
from .serializers import compat, yamlserializer, jsonserializer
|
|
from .serialize import serialize
|
|
from . import request
|
|
from .stubs.compat import get_httpmessage
|
|
|
|
# Use the libYAML versions if possible
|
|
try:
|
|
from yaml import CLoader as Loader
|
|
except ImportError:
|
|
from yaml import Loader
|
|
|
|
def preprocess_yaml(cassette):
|
|
return cassette.replace(' !!python/object:vcr.request.Request', '').replace('!!python/object/apply:__builtin__.frozenset', '')
|
|
|
|
|
|
PARTS = [
|
|
'protocol',
|
|
'host',
|
|
'port',
|
|
'path',
|
|
]
|
|
|
|
|
|
def build_uri(**parts):
|
|
port = parts['port']
|
|
scheme = parts['protocol']
|
|
default_port = {'https': 433, 'http': 80}[scheme]
|
|
parts['port'] = ':{0}'.format(port) if port != default_port else ''
|
|
return "{protocol}://{host}{port}{path}".format(**parts)
|
|
|
|
|
|
def _migrate(data):
|
|
interactions = []
|
|
for item in data:
|
|
req = item['request']
|
|
res = item['response']
|
|
uri = dict((k, req.pop(k)) for k in PARTS)
|
|
req['uri'] = build_uri(**uri)
|
|
# convert headers to dict of lists
|
|
headers = req['headers']
|
|
for k in headers:
|
|
headers[k] = [headers[k]]
|
|
response_headers = {}
|
|
for k, v in get_httpmessage(b"".join(h.encode('utf-8') for h in res['headers'])).items():
|
|
response_headers.setdefault(k, [])
|
|
response_headers[k].append(v)
|
|
res['headers'] = response_headers
|
|
interactions.append({'request':req, 'response': res})
|
|
return {
|
|
'requests': [request.Request._from_dict(i['request']) for i in interactions],
|
|
'responses': [i['response'] for i in interactions],
|
|
}
|
|
|
|
def migrate_json(in_fp, out_fp):
|
|
data = json.load(in_fp)
|
|
interactions = _migrate(data)
|
|
out_fp.write(serialize(interactions, jsonserializer))
|
|
|
|
|
|
def _list_of_tuples_to_dict(fs):
|
|
return dict((k, v) for k, v in fs[0])
|
|
|
|
|
|
def migrate_yml(in_fp, out_fp):
|
|
data = yaml.load(preprocess_yaml(in_fp.read()), Loader=Loader)
|
|
for i in range(len(data)):
|
|
data[i]['request']['headers'] = _list_of_tuples_to_dict(data[i]['request']['headers'])
|
|
interactions = _migrate(data)
|
|
out_fp.write(serialize(interactions, yamlserializer))
|
|
|
|
|
|
def migrate(file_path, migration_fn):
|
|
# because we assume that original files can be reverted
|
|
# we will try to copy the content. (os.rename not needed)
|
|
with tempfile.TemporaryFile(mode='w+') as out_fp:
|
|
with open(file_path, 'r') as in_fp:
|
|
migration_fn(in_fp, out_fp)
|
|
with open(file_path, 'w') as in_fp:
|
|
out_fp.seek(0)
|
|
shutil.copyfileobj(out_fp, in_fp)
|
|
|
|
|
|
def try_migrate(path):
|
|
if path.endswith('.json'):
|
|
migrate(path, migrate_json)
|
|
return True
|
|
elif path.endswith('.yaml') or path.endswith('.yml'):
|
|
migrate(path, migrate_yml)
|
|
return True
|
|
return False
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) != 2:
|
|
raise SystemExit("Please provide path to cassettes directory or file. "
|
|
"Usage: python -m vcr.migration PATH")
|
|
|
|
path = sys.argv[1]
|
|
if not os.path.isabs(path):
|
|
path = os.path.abspath(path)
|
|
files = [path]
|
|
if os.path.isdir(path):
|
|
files = (os.path.join(root, name)
|
|
for (root, dirs, files) in os.walk(path)
|
|
for name in files)
|
|
for file_path in files:
|
|
migrated = try_migrate(file_path)
|
|
status = 'OK' if migrated else 'FAIL'
|
|
sys.stderr.write("[{0}] {1}\n".format(status, file_path))
|
|
sys.stderr.write("Done.\n")
|
|
|
|
if __name__ == '__main__':
|
|
main()
|