Merge pull request #60 from firstbatchxyz/erhant/allow-testnet

added `DKN_NETWORK` env var

Author: erhant

Cargo.lock

@@ -1,7 +1,7 @@
 [package]
 name = "dkn-compute-launcher"
 edition = "2021"
-version = "0.1.17"
+version = "0.1.18"
 license = "Apache-2.0"
 readme = "README.md"
 description = "Dria Compute Node Launcher"

Cargo.toml

@@ -1,9 +1,14 @@
 use colored::Colorize;
 use eyre::Context;
 
-use crate::utils::{DriaEnv, LAUNCHER_USER_AGENT};
+use crate::utils::{get_network_env, DriaEnv, LAUNCHER_USER_AGENT};
 
-const POINTS_API_BASE_URL: &str = "https://mainnet.dkn.dria.co/points/v0/total/node/";
+#[inline]
+fn get_points_api_url(address: &str) -> String {
+    let network = get_network_env();
+    let address = address.trim_start_matches("0x");
+    format!("https://{network}.dkn.dria.co/points/v0/total/node/0x{address}")
+}
 
 #[derive(Debug, serde::Deserialize)]
 pub struct PointsRes {
@@ -43,11 +48,7 @@ pub async fn show_points() -> eyre::Result<()> {
 }
 
 async fn get_points(address: &str) -> eyre::Result<PointsRes> {
-    let url = format!(
-        "{}/0x{}",
-        POINTS_API_BASE_URL,
-        address.trim_start_matches("0x")
-    );
+    let url = get_points_api_url(address);
 
     let client = reqwest::Client::builder()
         .user_agent(LAUNCHER_USER_AGENT)

src/commands/points.rs

@@ -9,9 +9,6 @@ use crate::utils::{referrals::*, DriaEnv, Selectable};
 pub async fn handle_referrals() -> eyre::Result<()> {
     // ensure system is healthy
     let client = ReferralsClient::default();
-    if !client.healthcheck().await {
-        eyre::bail!("Referrals API is offline.");
-    }
 
     // get wallet secret from env
     let mut dria_env = DriaEnv::new_from_env();

src/commands/referrals.rs

@@ -48,3 +48,17 @@ pub const PROGRESS_BAR_CHARS: &str = "=>-";
 /// `UserAgent` header value for the launcher, used for HTTP requests.
 pub const LAUNCHER_USER_AGENT: &str =
     concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
+
+/// Returns the network type based on the `DKN_NETWORK` environment variable.
+///
+/// This is usually not used at all by the user, but rather used in testing and development.
+#[inline(always)]
+pub fn get_network_env() -> String {
+    std::env::var("DKN_NETWORK")
+        .map(|s| match s.as_str() {
+            // only accept `testnet` as a valid network, otherwise default to `mainnet`
+            "testnet" => s,
+            _ => "mainnet".to_string(),
+        })
+        .unwrap_or_else(|_| "mainnet".to_string())
+}

src/utils/mod.rs

@@ -1,11 +1,14 @@
+use crate::utils::{crypto::eip191_hash, get_network_env};
 use eyre::{Context, Result};
 use libsecp256k1::SecretKey;
 
-use crate::utils::crypto::eip191_hash;
-
 use super::LAUNCHER_USER_AGENT;
 
-const REFERRALS_API_BASE_URL: &str = "https://mainnet.dkn.dria.co/referrals/v0/";
+#[inline]
+fn get_referrals_api_base_url() -> String {
+    let network = get_network_env();
+    format!("https://{network}.dkn.dria.co/referral/v0")
+}
 
 pub struct ReferralsClient {
     base_url: String,
@@ -14,30 +17,18 @@ pub struct ReferralsClient {
 
 impl Default for ReferralsClient {
     fn default() -> Self {
-        Self::new(REFERRALS_API_BASE_URL)
+        Self::new(get_referrals_api_base_url())
     }
 }
 
 impl ReferralsClient {
-    pub fn new(base_url: &str) -> Self {
+    pub fn new(base_url: String) -> Self {
         let client = reqwest::Client::builder()
             .user_agent(LAUNCHER_USER_AGENT)
             .build()
             .expect("could not create reqwest client");
 
-        Self {
-            base_url: base_url.to_string(),
-            client,
-        }
-    }
-
-    /// Simple healthcheck for the referrals API.
-    pub async fn healthcheck(&self) -> bool {
-        self.client
-            .get(format!("{}/health", self.base_url))
-            .send()
-            .await
-            .is_ok_and(|r| r.status().is_success())
+        Self { base_url, client }
     }
 
     /// Returns a list of addresses of the users referred by the given `address`.
@@ -193,13 +184,6 @@ impl ReferralsClient {
 mod tests {
     use super::*;
 
-    #[tokio::test]
-    #[ignore]
-    async fn test_health() {
-        let ok = ReferralsClient::default().healthcheck().await;
-        assert!(ok);
-    }
-
     #[tokio::test]
     #[ignore]
     async fn test_get_referrals() {