Skip to content
This repository was archived by the owner on Mar 13, 2022. It is now read-only.

Commit 0b1ac80

Browse files
zshihangzshihang
committed
allow incluster to accept pass-in config
1 parent bf5c599 commit 0b1ac80

File tree

2 files changed

+85
-73
lines changed

2 files changed

+85
-73
lines changed

Diff for: ‎config/incluster_config.py

+45-35
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
import os
1615
import datetime
16+
import os
1717

1818
from kubernetes.client import Configuration
1919

@@ -35,30 +35,39 @@ def _join_host_port(host, port):
3535

3636

3737
class InClusterConfigLoader(object):
38-
39-
def __init__(self, token_filename,
40-
cert_filename, environ=os.environ):
38+
def __init__(self,
39+
token_filename,
40+
cert_filename,
41+
try_refresh_token=True,
42+
environ=os.environ):
4143
self._token_filename = token_filename
4244
self._cert_filename = cert_filename
4345
self._environ = environ
46+
self._try_refresh_token = try_refresh_token
4447
self._token_refresh_period = datetime.timedelta(minutes=1)
4548

46-
def load_and_set(self, refresh_token=True):
49+
def load_and_set(self, client_configuration=None):
50+
try_set_default = False
51+
if client_configuration is None:
52+
client_configuration = type.__call__(Configuration)
53+
try_set_default = True
4754
self._load_config()
48-
self._set_config(refresh_token=refresh_token)
55+
self._set_config(client_configuration)
56+
if try_set_default:
57+
Configuration.set_default(client_configuration)
4958

5059
def _load_config(self):
51-
if (SERVICE_HOST_ENV_NAME not in self._environ or
52-
SERVICE_PORT_ENV_NAME not in self._environ):
60+
if (SERVICE_HOST_ENV_NAME not in self._environ
61+
or SERVICE_PORT_ENV_NAME not in self._environ):
5362
raise ConfigException("Service host/port is not set.")
5463

55-
if (not self._environ[SERVICE_HOST_ENV_NAME] or
56-
not self._environ[SERVICE_PORT_ENV_NAME]):
64+
if (not self._environ[SERVICE_HOST_ENV_NAME]
65+
or not self._environ[SERVICE_PORT_ENV_NAME]):
5766
raise ConfigException("Service host/port is set but empty.")
5867

59-
self.host = (
60-
"https://" + _join_host_port(self._environ[SERVICE_HOST_ENV_NAME],
61-
self._environ[SERVICE_PORT_ENV_NAME]))
68+
self.host = ("https://" +
69+
_join_host_port(self._environ[SERVICE_HOST_ENV_NAME],
70+
self._environ[SERVICE_PORT_ENV_NAME]))
6271

6372
if not os.path.isfile(self._token_filename):
6473
raise ConfigException("Service token file does not exists.")
@@ -75,37 +84,38 @@ def _load_config(self):
7584

7685
self.ssl_ca_cert = self._cert_filename
7786

78-
def _set_config(self, refresh_token):
79-
configuration = Configuration()
80-
configuration.host = self.host
81-
configuration.ssl_ca_cert = self.ssl_ca_cert
82-
configuration.api_key['authorization'] = "bearer " + self.token
83-
Configuration.set_default(configuration)
84-
if not refresh_token:
87+
def _set_config(self, client_configuration):
88+
client_configuration.host = self.host
89+
client_configuration.ssl_ca_cert = self.ssl_ca_cert
90+
if self.token is not None:
91+
client_configuration.api_key['authorization'] = self.token
92+
if not self._try_refresh_token:
8593
return
86-
def wrap(f):
87-
in_cluster_config = self
88-
def wrapped(self, identifier):
89-
if identifier == 'authorization' and identifier in self.api_key and in_cluster_config.token_expires_at <= datetime.datetime.now():
90-
in_cluster_config._read_token_file()
91-
self.api_key[identifier] = "bearer " + in_cluster_config.token
92-
return f(self, identifier)
93-
return wrapped
94-
Configuration.get_api_key_with_prefix = wrap(Configuration.get_api_key_with_prefix)
94+
95+
def load_token_from_file(*args):
96+
if self.token_expires_at <= datetime.datetime.now():
97+
self._read_token_file()
98+
return self.token
99+
100+
client_configuration.get_api_key_with_prefix = load_token_from_file
95101

96102
def _read_token_file(self):
97103
with open(self._token_filename) as f:
98-
self.token = f.read()
99-
self.token_expires_at = datetime.datetime.now() + self._token_refresh_period
100-
if not self.token:
104+
content = f.read()
105+
if not content:
101106
raise ConfigException("Token file exists but empty.")
107+
self.token = "bearer " + content
108+
self.token_expires_at = datetime.datetime.now(
109+
) + self._token_refresh_period
102110

103111

104-
def load_incluster_config(refresh_token=True):
112+
def load_incluster_config(client_configuration=None, try_refresh_token=True):
105113
"""
106114
Use the service account kubernetes gives to pods to connect to kubernetes
107115
cluster. It's intended for clients that expect to be running inside a pod
108116
running on kubernetes. It will raise an exception if called from a process
109117
not running in a kubernetes environment."""
110-
InClusterConfigLoader(token_filename=SERVICE_TOKEN_FILENAME,
111-
cert_filename=SERVICE_CERT_FILENAME).load_and_set(refresh_token=refresh_token)
118+
InClusterConfigLoader(
119+
token_filename=SERVICE_TOKEN_FILENAME,
120+
cert_filename=SERVICE_CERT_FILENAME,
121+
try_refresh_token=try_refresh_token).load_and_set(client_configuration)

Diff for: ‎config/incluster_config_test.py

+40-38
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,11 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import datetime
1516
import os
1617
import tempfile
17-
import unittest
18-
import datetime
1918
import time
19+
import unittest
2020

2121
from kubernetes.client import Configuration
2222

@@ -33,14 +33,17 @@
3333
_TEST_IPV6_HOST = "::1"
3434
_TEST_IPV6_HOST_PORT = "[::1]:80"
3535

36-
_TEST_ENVIRON = {SERVICE_HOST_ENV_NAME: _TEST_HOST,
37-
SERVICE_PORT_ENV_NAME: _TEST_PORT}
38-
_TEST_IPV6_ENVIRON = {SERVICE_HOST_ENV_NAME: _TEST_IPV6_HOST,
39-
SERVICE_PORT_ENV_NAME: _TEST_PORT}
36+
_TEST_ENVIRON = {
37+
SERVICE_HOST_ENV_NAME: _TEST_HOST,
38+
SERVICE_PORT_ENV_NAME: _TEST_PORT
39+
}
40+
_TEST_IPV6_ENVIRON = {
41+
SERVICE_HOST_ENV_NAME: _TEST_IPV6_HOST,
42+
SERVICE_PORT_ENV_NAME: _TEST_PORT
43+
}
4044

4145

4246
class InClusterConfigTest(unittest.TestCase):
43-
4447
def setUp(self):
4548
self._temp_files = []
4649

@@ -55,25 +58,18 @@ def _create_file_with_temp_content(self, content=""):
5558
os.close(handler)
5659
return name
5760

58-
def _overwrite_file_with_content(self, name, content=""):
59-
handler = os.open(name, os.O_RDWR)
60-
os.truncate(name, 0)
61-
os.write(handler, str.encode(content))
62-
os.close(handler)
63-
64-
def get_test_loader(
65-
self,
66-
token_filename=None,
67-
cert_filename=None,
68-
environ=_TEST_ENVIRON):
61+
def get_test_loader(self,
62+
token_filename=None,
63+
cert_filename=None,
64+
environ=_TEST_ENVIRON):
6965
if not token_filename:
7066
token_filename = self._create_file_with_temp_content(_TEST_TOKEN)
7167
if not cert_filename:
7268
cert_filename = self._create_file_with_temp_content(_TEST_CERT)
73-
return InClusterConfigLoader(
74-
token_filename=token_filename,
75-
cert_filename=cert_filename,
76-
environ=environ)
69+
return InClusterConfigLoader(token_filename=token_filename,
70+
cert_filename=cert_filename,
71+
try_refresh_token=True,
72+
environ=environ)
7773

7874
def test_join_host_port(self):
7975
self.assertEqual(_TEST_HOST_PORT,
@@ -87,25 +83,29 @@ def test_load_config(self):
8783
loader._load_config()
8884
self.assertEqual("https://" + _TEST_HOST_PORT, loader.host)
8985
self.assertEqual(cert_filename, loader.ssl_ca_cert)
90-
self.assertEqual(_TEST_TOKEN, loader.token)
86+
self.assertEqual('bearer ' + _TEST_TOKEN, loader.token)
9187

9288
def test_refresh_token(self):
9389
loader = self.get_test_loader()
94-
loader._token_refresh_period = datetime.timedelta(seconds=5)
95-
loader.load_and_set()
9690
config = Configuration()
91+
loader.load_and_set(config)
9792

98-
self.assertEqual('bearer '+_TEST_TOKEN, config.get_api_key_with_prefix('authorization'))
99-
self.assertEqual(_TEST_TOKEN, loader.token)
93+
self.assertEqual('bearer ' + _TEST_TOKEN,
94+
config.get_api_key_with_prefix('authorization'))
95+
self.assertEqual('bearer ' + _TEST_TOKEN, loader.token)
10096
self.assertIsNotNone(loader.token_expires_at)
10197

10298
old_token = loader.token
10399
old_token_expires_at = loader.token_expires_at
104-
self._overwrite_file_with_content(loader._token_filename, _TEST_NEW_TOKEN)
105-
time.sleep(5)
106-
107-
self.assertEqual('bearer '+_TEST_NEW_TOKEN, config.get_api_key_with_prefix('authorization'))
108-
self.assertEqual(_TEST_NEW_TOKEN, loader.token)
100+
loader._token_filename = self._create_file_with_temp_content(
101+
_TEST_NEW_TOKEN)
102+
self.assertEqual('bearer ' + _TEST_TOKEN,
103+
config.get_api_key_with_prefix('authorization'))
104+
105+
loader.token_expires_at = datetime.datetime.now()
106+
self.assertEqual('bearer ' + _TEST_NEW_TOKEN,
107+
config.get_api_key_with_prefix('authorization'))
108+
self.assertEqual('bearer ' + _TEST_NEW_TOKEN, loader.token)
109109
self.assertGreater(loader.token_expires_at, old_token_expires_at)
110110

111111
def _should_fail_load(self, config_loader, reason):
@@ -122,9 +122,10 @@ def test_no_port(self):
122122
self._should_fail_load(loader, "no port specified")
123123

124124
def test_empty_port(self):
125-
loader = self.get_test_loader(
126-
environ={SERVICE_HOST_ENV_NAME: _TEST_HOST,
127-
SERVICE_PORT_ENV_NAME: ""})
125+
loader = self.get_test_loader(environ={
126+
SERVICE_HOST_ENV_NAME: _TEST_HOST,
127+
SERVICE_PORT_ENV_NAME: ""
128+
})
128129
self._should_fail_load(loader, "empty port specified")
129130

130131
def test_no_host(self):
@@ -133,9 +134,10 @@ def test_no_host(self):
133134
self._should_fail_load(loader, "no host specified")
134135

135136
def test_empty_host(self):
136-
loader = self.get_test_loader(
137-
environ={SERVICE_HOST_ENV_NAME: "",
138-
SERVICE_PORT_ENV_NAME: _TEST_PORT})
137+
loader = self.get_test_loader(environ={
138+
SERVICE_HOST_ENV_NAME: "",
139+
SERVICE_PORT_ENV_NAME: _TEST_PORT
140+
})
139141
self._should_fail_load(loader, "empty host specified")
140142

141143
def test_no_cert_file(self):

0 commit comments

Comments
 (0)