|
| 1 | +# run with "locust -f locust_files/load_testing.py" |
| 2 | +# then set number of users and spawn rate in web interface |
| 3 | +# e.g. test with 100 users, and 15 users/sec |
| 4 | +# web interface: http://0.0.0.0:8089/ |
| 5 | + |
| 6 | +import datetime |
| 7 | +import json |
| 8 | +import random |
| 9 | +from uuid import uuid4 |
| 10 | + |
| 11 | +from locust import HttpUser, between, task |
| 12 | + |
| 13 | +from mapswipe_workers.definitions import logger |
| 14 | +from mapswipe_workers.utils import user_management |
| 15 | + |
| 16 | + |
| 17 | +class MapSwipeUser(HttpUser): |
| 18 | + # assuming that users need between 30 sec and 120 sec to map a group |
| 19 | + wait_time = between(30, 120) |
| 20 | + |
| 21 | + def set_up_user(self): |
| 22 | + # check if is already signed in |
| 23 | + if self.signed_in_user is None: |
| 24 | + logger.info("user is not signed in. Will create a new user.") |
| 25 | + # create user if not exists |
| 26 | + user = user_management.create_user(self.email, self.username, self.password) |
| 27 | + self.user_id = user.uid |
| 28 | + |
| 29 | + # sign in user |
| 30 | + self.signed_in_user = user_management.sign_in_with_email_and_password( |
| 31 | + self.email, self.password |
| 32 | + ) |
| 33 | + logger.info("Created a new user.") |
| 34 | + else: |
| 35 | + logger.info("user is already signed in.") |
| 36 | + pass |
| 37 | + |
| 38 | + def create_mock_result(self, group): |
| 39 | + """Create a result object for a build area project. |
| 40 | +
|
| 41 | + The result values are generated randomly. |
| 42 | + """ |
| 43 | + start_time = datetime.datetime.utcnow().isoformat()[0:-3] + "Z" |
| 44 | + end_time = datetime.datetime.utcnow().isoformat()[0:-3] + "Z" |
| 45 | + |
| 46 | + x_min = int(group["xMin"]) |
| 47 | + x_max = int(group["xMax"]) |
| 48 | + y_min = int(group["yMin"]) |
| 49 | + y_max = int(group["yMax"]) |
| 50 | + |
| 51 | + results = {} |
| 52 | + for x in range(x_min, x_max): |
| 53 | + for y in range(y_min, y_max): |
| 54 | + task_id = f"18-{x}-{y}" |
| 55 | + result = random.choices([0, 1, 2, 3])[0] # no, yes, maybe, bad_imagery |
| 56 | + results[task_id] = result |
| 57 | + |
| 58 | + data = { |
| 59 | + "results": results, |
| 60 | + "startTime": start_time, |
| 61 | + "endTime": end_time, |
| 62 | + } |
| 63 | + return data |
| 64 | + |
| 65 | + def set_firebase_db(self, path, data, token=None): |
| 66 | + """Upload results to Firebase using REST api.""" |
| 67 | + request_ref = f"{path}.json?auth={token}" |
| 68 | + headers = {"content-type": "application/json; charset=UTF-8"} |
| 69 | + self.client.put( |
| 70 | + request_ref, headers=headers, data=json.dumps(data).encode("utf-8") |
| 71 | + ) |
| 72 | + logger.info(f"set data in firebase for {path}.json") |
| 73 | + |
| 74 | + @task |
| 75 | + def map_a_group(self): |
| 76 | + """Get a group from Firebase for this user and "map" it. |
| 77 | +
|
| 78 | + Make sure that this user has not worked on this group before. |
| 79 | + Get the group and create mock results. |
| 80 | + Upload results to Firebse. |
| 81 | + """ |
| 82 | + |
| 83 | + # get the groups that need to be mapped |
| 84 | + path = f"/v2/groups/{self.project_id}" |
| 85 | + # make sure to set '&' at the end of the string |
| 86 | + custom_arguments = 'orderBy="requiredCount"&limitToLast=15&' |
| 87 | + new_groups = user_management.get_firebase_db( |
| 88 | + path, custom_arguments, self.signed_in_user["idToken"] |
| 89 | + ) |
| 90 | + |
| 91 | + # get the groups the user has worked on already |
| 92 | + path = f"/v2/users/{self.user_id}/contributions/{self.project_id}" |
| 93 | + # make sure to set & at the end of the string |
| 94 | + custom_arguments = "shallow=True&" |
| 95 | + existing_groups = user_management.get_firebase_db( |
| 96 | + path, custom_arguments, self.signed_in_user["idToken"] |
| 97 | + ) |
| 98 | + |
| 99 | + # pick group for mapping |
| 100 | + # Get difference between new_groups and existing groups. |
| 101 | + # We should get the groups the user has not yet worked on. |
| 102 | + if existing_groups is None: |
| 103 | + next_group_id = random.choice(list(new_groups.keys())) |
| 104 | + else: |
| 105 | + existing_groups.pop( |
| 106 | + "taskContributionCount", None |
| 107 | + ) # need to remove this since it is no groupId |
| 108 | + remaining_group_ids = list( |
| 109 | + set(new_groups.keys()) - set(existing_groups.keys()) |
| 110 | + ) |
| 111 | + next_group_id = random.choice(remaining_group_ids) |
| 112 | + |
| 113 | + # get group object |
| 114 | + next_group = new_groups[next_group_id] |
| 115 | + |
| 116 | + # create mock result for this group |
| 117 | + result = self.create_mock_result(next_group) |
| 118 | + |
| 119 | + # upload results in firebase |
| 120 | + path = f"/v2/results/{self.project_id}/{next_group_id}/{self.user_id}" |
| 121 | + self.set_firebase_db(path, result, self.signed_in_user["idToken"]) |
| 122 | + |
| 123 | + def on_start(self): |
| 124 | + """Set up user and define project when user starts running.""" |
| 125 | + self.project_id = "-MYg8CEf2k1-RitN62X0" |
| 126 | + random_string = uuid4() |
| 127 | + self.email = f"test_{random_string}@mapswipe.org" |
| 128 | + self.username = f"test_{random_string}" |
| 129 | + self.password = "mapswipe" |
| 130 | + self.user_id = None |
| 131 | + self.signed_in_user = None |
| 132 | + self.set_up_user() |
0 commit comments