docker-python/tests/test_user_secrets.py at master · Technology-Development/docker-python · GitHub

History

218 lines (181 loc) · 9.43 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

import json

import os

import subprocess

import threading

import unittest

from http.server import BaseHTTPRequestHandler, HTTPServer

from test.support import EnvironmentVarGuard

from urllib.parse import urlparse

from datetime import datetime, timedelta

import mock

from google.auth.exceptions import DefaultCredentialsError

from google.cloud import bigquery

from kaggle_secrets import (GcpTarget, UserSecretsClient,

NotFoundError, ValidationError)

from kaggle_web_client import (_KAGGLE_URL_BASE_ENV_VAR_NAME,

_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME,

CredentialError, BackendError)

_TEST_JWT = 'test-secrets-key'

class UserSecretsHTTPHandler(BaseHTTPRequestHandler):

def set_request(self):

raise NotImplementedError()

def get_response(self):

raise NotImplementedError()

def do_HEAD(s):

s.send_response(200)

def do_POST(s):

s.set_request()

s.send_response(200)

s.send_header("Content-type", "application/json")

s.end_headers()

s.wfile.write(json.dumps(s.get_response()).encode("utf-8"))

class TestUserSecrets(unittest.TestCase):

SERVER_ADDRESS = urlparse(os.getenv(_KAGGLE_URL_BASE_ENV_VAR_NAME, default="http://127.0.0.1:8001"))

def _test_client(self, client_func, expected_path, expected_body, secret=None, success=True):

_request = {}

class AccessTokenHandler(UserSecretsHTTPHandler):

def set_request(self):

_request['path'] = self.path

content_len = int(self.headers.get('Content-Length'))

_request['body'] = json.loads(self.rfile.read(content_len))

_request['headers'] = self.headers

def get_response(self):

if success:

return {'result': {'secret': secret, 'secretType': 'refreshToken', 'secretProvider': 'google', 'expiresInSeconds': 3600}, 'wasSuccessful': "true"}

else:

return {'wasSuccessful': "false", 'errors': ['No user secrets exist for kernel']}

env = EnvironmentVarGuard()

env.set(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME, _TEST_JWT)

with env:

with HTTPServer((self.SERVER_ADDRESS.hostname, self.SERVER_ADDRESS.port), AccessTokenHandler) as httpd:

threading.Thread(target=httpd.serve_forever).start()

try:

client_func()

finally:

httpd.shutdown()

path, headers, body = _request['path'], _request['headers'], _request['body']

self.assertEqual(

path,

expected_path,

msg="Fake server did not receive the right request from the UserSecrets client.")

self.assertEqual(

body,

expected_body,

msg="Fake server did not receive the right body from the UserSecrets client.")

def test_no_token_fails(self):

env = EnvironmentVarGuard()

env.unset(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME)

with env:

with self.assertRaises(CredentialError):

client = UserSecretsClient()

def test_get_secret_succeeds(self):

secret = '12345'

def call_get_secret():

client = UserSecretsClient()

secret_response = client.get_secret("secret_label")

self.assertEqual(secret_response, secret)

self._test_client(call_get_secret,

'/requests/GetUserSecretByLabelRequest', {'Label': "secret_label"},

secret=secret)

def test_get_secret_handles_unsuccessful(self):

def call_get_secret():

client = UserSecretsClient()

with self.assertRaises(BackendError):

secret_response = client.get_secret("secret_label")

self._test_client(call_get_secret,

'/requests/GetUserSecretByLabelRequest', {'Label': "secret_label"},

success=False)

def test_get_secret_validates_label(self):

env = EnvironmentVarGuard()

env.set(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME, _TEST_JWT)

with env:

client = UserSecretsClient()

with self.assertRaises(ValidationError):

secret_response = client.get_secret("")

def test_get_gcloud_secret_succeeds(self):

secret = '{"client_id":"gcloud","type":"authorized_user"}'

def call_get_secret():

client = UserSecretsClient()

secret_response = client.get_gcloud_credential()

self.assertEqual(secret_response, secret)

self._test_client(call_get_secret,

'/requests/GetUserSecretByLabelRequest', {'Label': "__gcloud_sdk_auth__"},

secret=secret)

def test_get_gcloud_secret_handles_unsuccessful(self):

def call_get_secret():

client = UserSecretsClient()

with self.assertRaises(NotFoundError):

secret_response = client.get_gcloud_credential()

self._test_client(call_get_secret,

'/requests/GetUserSecretByLabelRequest', {'Label': "__gcloud_sdk_auth__"},

success=False)

def test_set_gcloud_credentials_succeeds(self):

secret = '{"client_id":"gcloud","type":"authorized_user","refresh_token":"refresh_token"}'

project = 'foo'

account = 'bar'

def get_gcloud_config_value(field):

result = subprocess.run(['gcloud', 'config', 'get-value', field], capture_output=True)

result.check_returncode()

return result.stdout.strip().decode('ascii')

def test_fn():

client = UserSecretsClient()

client.set_gcloud_credentials(project=project, account=account)

self.assertEqual(project, os.environ['GOOGLE_CLOUD_PROJECT'])

self.assertEqual(project, get_gcloud_config_value('project'))

self.assertEqual(account, os.environ['GOOGLE_ACCOUNT'])

self.assertEqual(account, get_gcloud_config_value('account'))

expected_creds_file = '/tmp/gcloud_credential.json'

self.assertEqual(expected_creds_file, os.environ['GOOGLE_APPLICATION_CREDENTIALS'])

self.assertEqual(expected_creds_file, get_gcloud_config_value('auth/credential_file_override'))

with open(expected_creds_file, 'r') as f:

self.assertEqual(secret, '\n'.join(f.readlines()))

self._test_client(test_fn, '/requests/GetUserSecretByLabelRequest', {'Label': "__gcloud_sdk_auth__"}, secret=secret)

@mock.patch('kaggle_secrets.datetime')

def test_get_access_token_succeeds(self, mock_dt):

secret = '12345'

now = datetime(1993, 4, 24)

mock_dt.utcnow = mock.Mock(return_value=now)

def call_get_bigquery_access_token():

client = UserSecretsClient()

secret_response = client.get_bigquery_access_token()