Add unit test for router_pick_directory_server_impl

4 files changed, 45 insertions, 4 deletions

diff --git a/src/or/networkstatus.c b/src/or/networkstatus.c
index 173c109d60..076a2c4116 100644
--- a/src/or/networkstatus.c
+++ b/src/or/networkstatus.c
@@ -1457,6 +1457,38 @@ networkstatus_copy_old_consensus_info(networkstatus_t *new_c,
   } SMARTLIST_FOREACH_JOIN_END(rs_old, rs_new);
 }
 
+#ifdef TOR_UNIT_TESTS
+/**Accept a <b>flavor</b> consensus <b>c</b> without any additional
+ * validation. This is exclusively for unit tests.
+ * We copy any ancillary information from a pre-existing consensus
+ * and then free the current one and replace it with the newly
+ * provided instance. Returns -1 on unrecognized flavor, 0 otherwise.
+ */
+int
+networkstatus_set_current_consensus_from_ns(networkstatus_t *c,
+                                            const char *flavor)
+{
+  int flav = networkstatus_parse_flavor_name(flavor);
+  switch (flav) {
+    case FLAV_NS:
+      if (current_ns_consensus) {
+        networkstatus_copy_old_consensus_info(c, current_ns_consensus);
+        networkstatus_vote_free(current_ns_consensus);
+      }
+      current_ns_consensus = c;
+      break;
+    case FLAV_MICRODESC:
+      if (current_md_consensus) {
+        networkstatus_copy_old_consensus_info(c, current_md_consensus);
+        networkstatus_vote_free(current_md_consensus);
+      }
+      current_md_consensus = c;
+      break;
+  }
+  return current_md_consensus ? 0 : -1;
+}
+#endif //TOR_UNIT_TESTS
+
 /** Try to replace the current cached v3 networkstatus with the one in
  * <b>consensus</b>.  If we don't have enough certificates to validate it,
  * store it in consensus_waiting_for_certs and launch a certificate fetch.
diff --git a/src/or/networkstatus.h b/src/or/networkstatus.h
index 4cb33c3fc0..4eab4d83f8 100644
--- a/src/or/networkstatus.h
+++ b/src/or/networkstatus.h
@@ -114,6 +114,10 @@ int networkstatus_get_weight_scale_param(networkstatus_t *ns);
 
 #ifdef NETWORKSTATUS_PRIVATE
 STATIC void vote_routerstatus_free(vote_routerstatus_t *rs);
+#ifdef TOR_UNIT_TESTS
+STATIC int networkstatus_set_current_consensus_from_ns(networkstatus_t *c,
+                                                const char *flavor);
+#endif // TOR_UNIT_TESTS
 #endif
 
 #endif
diff --git a/src/or/routerlist.c b/src/or/routerlist.c
index df6d797e0d..af6fd747bf 100644
--- a/src/or/routerlist.c
+++ b/src/or/routerlist.c
@@ -67,8 +67,6 @@ typedef struct cert_list_t cert_list_t;
 static int compute_weighted_bandwidths(const smartlist_t *sl,
                                        bandwidth_weight_rule_t rule,
                                        u64_dbl_t **bandwidths_out);
-static const routerstatus_t *router_pick_directory_server_impl(
-                              dirinfo_type_t auth, int flags, int *n_busy_out);
 static const routerstatus_t *router_pick_trusteddirserver_impl(
                 const smartlist_t *sourcelist, dirinfo_type_t auth,
                 int flags, int *n_busy_out);
@@ -1472,7 +1470,7 @@ router_pick_dirserver_generic(smartlist_t *sourcelist,
  * directories that we excluded for no other reason than
  * PDS_NO_EXISTING_SERVERDESC_FETCH or PDS_NO_EXISTING_MICRODESC_FETCH.
  */
-static const routerstatus_t *
+STATIC const routerstatus_t *
 router_pick_directory_server_impl(dirinfo_type_t type, int flags,
                                   int *n_busy_out)
 {
@@ -3238,7 +3236,11 @@ routerlist_reparse_old(routerlist_t *rl, signed_descriptor_t *sd)
   return ri;
 }
 
-/** Free all memory held by the routerlist module. */
+/** Free all memory held by the routerlist module.
+ * Note: Calling routerlist_free_all() should always be paired with
+ * a call to nodelist_free_all(). These should only be called during
+ * cleanup.
+ */
 void
 routerlist_free_all(void)
 {
diff --git a/src/or/routerlist.h b/src/or/routerlist.h
index 339e34ae03..dd88aeb179 100644
--- a/src/or/routerlist.h
+++ b/src/or/routerlist.h
@@ -233,6 +233,9 @@ STATIC int choose_array_element_by_weight(const u64_dbl_t *entries,
                                           int n_entries);
 STATIC void scale_array_elements_to_u64(u64_dbl_t *entries, int n_entries,
                                         uint64_t *total_out);
+STATIC const routerstatus_t *router_pick_directory_server_impl(
+                                           dirinfo_type_t auth, int flags,
+                                           int *n_busy_out);
 
 MOCK_DECL(int, router_descriptor_is_older_than, (const routerinfo_t *router,
                                                  int seconds));