mirror of
https://github.com/mozilla-services/syncstorage-rs.git
synced 2026-01-21 00:12:14 +01:00
8995db6268
Some checks are pending
Glean probe-scraper / glean-probe-scraper (push) Waiting to run
feat: ruff for python lint and format
222 lines
8.6 KiB
Python
222 lines
8.6 KiB
Python
# This Source Code Form is subject to the terms of the Mozilla Public
|
|
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
|
|
# You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
|
|
import json
|
|
import os
|
|
import unittest
|
|
import uuid
|
|
|
|
from add_node import main as add_node_script
|
|
from allocate_user import main as allocate_user_script
|
|
from count_users import main as count_users_script
|
|
from database import Database
|
|
from remove_node import main as remove_node_script
|
|
from unassign_node import main as unassign_node_script
|
|
from update_node import main as update_node_script
|
|
from util import get_timestamp
|
|
|
|
|
|
class TestScripts(unittest.TestCase):
|
|
NODE_ID = 800
|
|
NODE_URL = "https://node1"
|
|
|
|
def setUp(self):
|
|
self.database = Database()
|
|
|
|
# Start each test with a blank slate.
|
|
cursor = self.database._execute_sql("DELETE FROM users")
|
|
cursor.close()
|
|
|
|
cursor = self.database._execute_sql("DELETE FROM nodes")
|
|
cursor.close()
|
|
|
|
cursor = self.database._execute_sql("DELETE FROM services")
|
|
cursor.close()
|
|
|
|
# Add a service
|
|
self.database.add_service("sync-1.5", r"{node}/1.5/{uid}")
|
|
|
|
# Ensure we have a node with enough capacity to run the tests.
|
|
self.database.add_node(self.NODE_URL, 100, id=self.NODE_ID)
|
|
|
|
def tearDown(self):
|
|
# And clean up at the end, for good measure.
|
|
cursor = self.database._execute_sql("DELETE FROM users")
|
|
cursor.close()
|
|
|
|
cursor = self.database._execute_sql("DELETE FROM nodes")
|
|
cursor.close()
|
|
|
|
cursor = self.database._execute_sql("DELETE FROM services")
|
|
cursor.close()
|
|
|
|
self.database.close()
|
|
|
|
def test_add_node(self):
|
|
add_node_script(args=["--current-load", "9", "test_node", "100"])
|
|
res = self.database.get_node("test_node")
|
|
# The node should have the expected attributes
|
|
self.assertEqual(res.capacity, 100)
|
|
self.assertEqual(res.available, 10)
|
|
self.assertEqual(res.current_load, 9)
|
|
self.assertEqual(res.downed, 0)
|
|
self.assertEqual(res.backoff, 0)
|
|
self.assertEqual(res.service, self.database.service_id)
|
|
|
|
def test_add_node_with_explicit_available(self):
|
|
args = ["--current-load", "9", "--available", "5", "test_node", "100"]
|
|
add_node_script(args=args)
|
|
res = self.database.get_node("test_node")
|
|
# The node should have the expected attributes
|
|
self.assertEqual(res.capacity, 100)
|
|
self.assertEqual(res.available, 5)
|
|
self.assertEqual(res.current_load, 9)
|
|
self.assertEqual(res.downed, 0)
|
|
self.assertEqual(res.backoff, 0)
|
|
self.assertEqual(res.service, self.database.service_id)
|
|
|
|
def test_add_downed_node(self):
|
|
add_node_script(args=["--downed", "test_node", "100"])
|
|
res = self.database.get_node("test_node")
|
|
# The node should have the expected attributes
|
|
self.assertEqual(res.capacity, 100)
|
|
self.assertEqual(res.available, 10)
|
|
self.assertEqual(res.current_load, 0)
|
|
self.assertEqual(res.downed, 1)
|
|
self.assertEqual(res.backoff, 0)
|
|
self.assertEqual(res.service, self.database.service_id)
|
|
|
|
def test_add_backoff_node(self):
|
|
add_node_script(args=["--backoff", "test_node", "100"])
|
|
res = self.database.get_node("test_node")
|
|
# The node should have the expected attributes
|
|
self.assertEqual(res.capacity, 100)
|
|
self.assertEqual(res.available, 10)
|
|
self.assertEqual(res.current_load, 0)
|
|
self.assertEqual(res.downed, 0)
|
|
self.assertEqual(res.backoff, 1)
|
|
self.assertEqual(res.service, self.database.service_id)
|
|
|
|
def test_allocate_user_user_already_exists(self):
|
|
email = "test@test.com"
|
|
self.database.allocate_user(email)
|
|
node = "https://node2"
|
|
self.database.add_node(node, 100)
|
|
allocate_user_script(args=[email, node])
|
|
user = self.database.get_user(email)
|
|
# The user should be assigned to the given node
|
|
self.assertEqual(user["node"], node)
|
|
# Another user should not have been created
|
|
count = self.database.count_users()
|
|
self.assertEqual(count, 1)
|
|
|
|
def test_allocate_user_given_node(self):
|
|
email = "test@test.com"
|
|
node = "https://node2"
|
|
self.database.add_node(node, 100)
|
|
allocate_user_script(args=[email, node])
|
|
user = self.database.get_user(email)
|
|
# A new user should be created and assigned to the given node
|
|
self.assertEqual(user["node"], node)
|
|
|
|
def test_allocate_user_not_given_node(self):
|
|
email = "test@test.com"
|
|
self.database.add_node("https://node2", 100, current_load=10)
|
|
self.database.add_node("https://node3", 100, current_load=20)
|
|
self.database.add_node("https://node4", 100, current_load=30)
|
|
allocate_user_script(args=[email])
|
|
user = self.database.get_user(email)
|
|
# The user should be assigned to the least-loaded node
|
|
self.assertEqual(user["node"], "https://node1")
|
|
|
|
def test_count_users(self):
|
|
self.database.allocate_user("test1@test.com")
|
|
self.database.allocate_user("test2@test.com")
|
|
self.database.allocate_user("test3@test.com")
|
|
|
|
timestamp = get_timestamp()
|
|
filename = "/tmp/" + str(uuid.uuid4())
|
|
try:
|
|
count_users_script(
|
|
args=["--output", filename, "--timestamp", str(timestamp)]
|
|
)
|
|
|
|
with open(filename) as f:
|
|
info = json.loads(f.readline())
|
|
self.assertEqual(info["total_users"], 3)
|
|
self.assertEqual(info["op"], "sync_count_users")
|
|
finally:
|
|
os.remove(filename)
|
|
|
|
filename = "/tmp/" + str(uuid.uuid4())
|
|
try:
|
|
args = ["--output", filename, "--timestamp", str(timestamp - 10000)]
|
|
count_users_script(args=args)
|
|
|
|
with open(filename) as f:
|
|
info = json.loads(f.readline())
|
|
self.assertEqual(info["total_users"], 0)
|
|
self.assertEqual(info["op"], "sync_count_users")
|
|
finally:
|
|
os.remove(filename)
|
|
|
|
def test_remove_node(self):
|
|
self.database.add_node("https://node2", 100)
|
|
self.database.allocate_user("test1@test.com", node="https://node2")
|
|
self.database.allocate_user("test2@test.com", node=self.NODE_URL)
|
|
self.database.allocate_user("test3@test.com", node=self.NODE_URL)
|
|
|
|
remove_node_script(args=["https://node2"])
|
|
|
|
# The node should have been removed from the database
|
|
args = ["https://node2"]
|
|
self.assertRaises(ValueError, self.database.get_node_id, *args)
|
|
# The first user should have been assigned to a new node
|
|
user = self.database.get_user("test1@test.com")
|
|
self.assertEqual(user["node"], self.NODE_URL)
|
|
# The second and third users should still be on the first node
|
|
user = self.database.get_user("test2@test.com")
|
|
self.assertEqual(user["node"], self.NODE_URL)
|
|
user = self.database.get_user("test3@test.com")
|
|
self.assertEqual(user["node"], self.NODE_URL)
|
|
|
|
def test_unassign_node(self):
|
|
self.database.add_node("https://node2", 100)
|
|
self.database.allocate_user("test1@test.com", node="https://node2")
|
|
self.database.allocate_user("test2@test.com", node="https://node2")
|
|
self.database.allocate_user("test3@test.com", node=self.NODE_URL)
|
|
|
|
unassign_node_script(args=["https://node2"])
|
|
self.database.remove_node("https://node2")
|
|
# All of the users should now be assigned to the first node
|
|
user = self.database.get_user("test1@test.com")
|
|
self.assertEqual(user["node"], self.NODE_URL)
|
|
user = self.database.get_user("test2@test.com")
|
|
self.assertEqual(user["node"], self.NODE_URL)
|
|
user = self.database.get_user("test3@test.com")
|
|
self.assertEqual(user["node"], self.NODE_URL)
|
|
|
|
def test_update_node(self):
|
|
self.database.add_node("https://node2", 100)
|
|
update_node_script(
|
|
args=[
|
|
"--capacity",
|
|
"150",
|
|
"--available",
|
|
"125",
|
|
"--current-load",
|
|
"25",
|
|
"--downed",
|
|
"--backoff",
|
|
"https://node2",
|
|
]
|
|
)
|
|
node = self.database.get_node("https://node2")
|
|
# Ensure the node has the expected attributes
|
|
self.assertEqual(node["capacity"], 150)
|
|
self.assertEqual(node["available"], 125)
|
|
self.assertEqual(node["current_load"], 25)
|
|
self.assertEqual(node["downed"], 1)
|
|
self.assertEqual(node["backoff"], 1)
|