import hashlib
import pytest

from mock import patch
from flask import url_for
from playhouse.test_utils import assert_query_count

from app import instance_keys, app as realapp
from data import model
from data.cache import InMemoryDataModelCache
from data.database import ImageStorageLocation
from endpoints.test.shared import conduct_call
from util.security.registry_jwt import generate_bearer_token, build_context_and_subject
from test.fixtures import *

@pytest.mark.parametrize('method, endpoint', [
  ('GET', 'download_blob'),
  ('HEAD', 'check_blob_exists'),
])
def test_blob_caching(method, endpoint, client, app):
  digest = 'sha256:' + hashlib.sha256("a").hexdigest()
  location = ImageStorageLocation.get(name='local_us')
  model.blob.store_blob_record_and_temp_link('devtable', 'simple', digest, location, 1, 10000000)

  params = {
    'repository': 'devtable/simple',
    'digest': digest,
  }

  user = model.user.get_user('devtable')
  access = [{
    'type': 'repository',
    'name': 'devtable/simple',
    'actions': ['pull'],
  }]

  context, subject = build_context_and_subject(user=user)
  token = generate_bearer_token(realapp.config['SERVER_HOSTNAME'], subject, context, access, 600,
                                instance_keys)

  headers = {
    'Authorization': 'Bearer %s' % token,
  }

  # Run without caching to make sure the request works. This also preloads some of
  # our global model caches.
  conduct_call(client, 'v2.' + endpoint, url_for, method, params, expected_code=200,
               headers=headers)

  with patch('endpoints.v2.blob.model_cache', InMemoryDataModelCache()):
    # First request should make a DB query to retrieve the blob.
    with assert_query_count(1):
      conduct_call(client, 'v2.' + endpoint, url_for, method, params, expected_code=200,
                   headers=headers)

    # Subsequent requests should use the cached blob.
    with assert_query_count(0):
      conduct_call(client, 'v2.' + endpoint, url_for, method, params, expected_code=200,
                   headers=headers)