def authorized_client(requests_mock, api_key_file, fs):
fs.create_dir(os.path.dirname(api_key_file))
assert not os.path.exists(api_key_file), "api key file shouldn't exist yet"
json_data = {
"id": 1234,
"api_key": test_api_key,
"username": "******",
"ssh_key": "ssh-rsa AAAA..."
}
requests_mock.put(api_base_url + "/users/current/", json=json_data)
new_client = VastClient(api_key_file=api_key_file)
assert type(new_client) is VastClient, "Should be a VastClient object."
assert new_client.api_key is None, "Shouldn't have an api_key yet."
retVal = new_client.authenticate('john_doe', 'abc123')
assert retVal is new_client, "Client.authenticate should return self."
assert new_client.api_key == test_api_key
assert requests_mock.last_request.json() == {
'username': '******',
'password': '******'
}
check_attrs(new_client, json_data)
assert os.path.exists(api_key_file), "Should have created api key file."
with open(api_key_file) as f:
assert f.read(
) == new_client.api_key, "Saved API key should match client.api_key."
return new_client
def test_default_api_key_file(fs):
client = VastClient()
expanded_path = os.path.expanduser(default_api_key_file)
assert client.api_key_file == expanded_path, \
"api_key_file should default to %s"%default_api_key_file
client = VastClient(api_key_file=None)
assert client.api_key_file is None, "client.api_key_file should still be None."
def test_api_key_env_var(requests_mock, monkeypatch):
monkeypatch.setenv('VAST_API_KEY', test_api_key)
client = VastClient()
assert client.api_key == test_api_key, "api_key should be set from VAST_API_KEY env var."
requests_mock.get(api_base_url + "/users/current/?api_key=%s" %
(client.api_key),
json=stubs.user_json)
client.authenticate()
check_attrs(client, stubs.user_json)
def test_create_ssh_key(fs):
client = VastClient()
assert client.ssh_key_dir is not None, "ssh_key_dir should be set to the default value."
fs.create_dir(client.ssh_key_dir)
key_name = "foobar"
public_key_file = client.create_ssh_key(key_name)
assert public_key_file == os.path.join(client.ssh_key_dir,
key_name + ".pub")
assert fs.exists(
public_key_file), "Public key file should exist on mocked filesystem."
assert fs.exists(public_key_file[:-4]
), "Private key file should exist on mocked filesystem."
def test_bad_login(requests_mock, fs):
requests_mock.put(api_base_url + "/users/current/",
status_code=401,
reason='Unauthorized')
client = VastClient(api_key_file=None)
assert client.api_key_file == None
with pytest.raises(Unauthorized):
retVal = client.authenticate('aFakeUser', 'badPassword')
assert retVal.api_key is None
assert requests_mock.last_request.json() == {
'username': '******',
'password': '******'
}
def vast_client(requests_mock, monkeypatch, fs):
fs.create_file("~/.vast_api_key", contents=test_api_key)
monkeypatch.setenv('VAST_API_KEY', test_api_key)
client = VastClient()
assert client.api_key == test_api_key, "Should be using test api key"
# Add attributes expected by vast.py commands.
vast_client.url = "https://vast.ai/api/v0"
vast_client.raw = True
return client
def test_authentication_with_api_key_file(requests_mock, authorized_client,
fs):
print(authorized_client.api_key_file,
os.path.exists(authorized_client.api_key_file))
assert fs.exists(
authorized_client.api_key_file), "api_key_file should exist."
json_data = {
"id": 1234,
"username": "******",
"ssh_key": "ssh-rsa AAAA..."
}
requests_mock.get(api_base_url + "/users/current/?api_key=%s" %
(authorized_client.api_key),
json=json_data)
new_client = VastClient(api_key_file=authorized_client.api_key_file)
assert new_client.api_key is not None, "new_client.api_key should be set."
assert new_client.api_key == authorized_client.api_key, \
"new_client.api_key should match authorized_client.api_key."
new_client.authenticate()
assert new_client.api_key == authorized_client.api_key
check_attrs(new_client, json_data)
def test_username_password_env_vars(requests_mock, monkeypatch, fs):
""" Tests authentication with username and password set in environment variables.
Params:
requests_mock: Network request mocker.
monkeypatch: Used for mocking env variables.
fs: pyfakefs plugin for mocking filesystem operations.
"""
monkeypatch.setenv('VAST_USERNAME', 'john_doe')
monkeypatch.setenv('VAST_PASSWORD', 'abc123')
client = VastClient()
expanded_path = os.path.expanduser(default_api_key_file)
#assert client.api_key_file == os.path.dirname(expanded_path), "api_key_file should be in home dir."
assert client.api_key is None, "api_key shouldn't be set yet."
requests_mock.put(api_base_url + "/users/current/", json=stubs.user_json)
fs.create_dir(os.path.dirname(expanded_path))
client.authenticate()
assert client.api_key_file == expanded_path
with open(client.api_key_file) as f:
assert f.read(
) == client.api_key, "Contents of %s should match client.api_key" % default_api_key_file
check_attrs(client, stubs.user_json)
assert client.api_key == test_api_key, "Retrieved api key should match test_api_key."
def test_login(requests_mock, fs):
assert not fs.exists(
"~/.vast_api_key"), "API key shouldn't be on mocked filesystem."
fs.create_dir(os.path.expanduser("~"))
client = VastClient()
assert client.api_key is None, "api_key is not set."
login_opts = LoginArgs()
_compare_requests(requests_mock,
vast.login,
login_opts,
client.authenticate,
response_json=stubs.user_json,
username=login_opts.username,
password=login_opts.password)
with open(client.api_key_file) as f:
assert f.read(
) == client.api_key, "API key should match the key saved to disk."