observation-tools 0.0.13

Observation Tools Client Library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use observation_tools::server_client::types::GetObservation;
use observation_tools::Client;
use observation_tools::ClientBuilder;
use observation_tools_server::auth::generate_api_key;
use observation_tools_server::auth::ApiKeySecret;

/// Test server wrapper that provides convenient client creation
pub struct TestServer {
  base_url: String,
  api_secret: Option<ApiKeySecret>,
  _handle: Option<tokio::task::JoinHandle<()>>,
}

impl TestServer {
  pub async fn new() -> Self {
    Self::new_internal(None).await
  }

  /// Create a new test server with API key authentication enabled
  #[allow(unused)]
  pub async fn new_with_auth() -> anyhow::Result<Self> {
    use rand::Rng;
    let mut rng = rand::thread_rng();
    let random1: u64 = rng.gen();
    let random2: u64 = rng.gen();
    let secret = format!("test-secret-key-{:016x}{:016x}", random1, random2);
    let api_secret = ApiKeySecret::new(&secret)?;
    Ok(Self::new_internal(Some(api_secret)).await)
  }

  async fn new_internal(api_secret: Option<ApiKeySecret>) -> Self {
    if let Ok(url) = std::env::var("SERVER_URL") {
      return Self {
        base_url: url,
        api_secret,
        _handle: None,
      };
    }

    let data_dir = tempfile::tempdir().expect("Failed to create temp dir");

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
      .await
      .expect("Failed to bind to random port");

    let addr = listener.local_addr().expect("Failed to get local address");

    let config = observation_tools_server::Config::new()
      .with_bind_addr(addr)
      .with_data_dir(data_dir.path().to_path_buf())
      .with_api_secret(api_secret.clone());

    let server = observation_tools_server::Server::new(config);

    let handle = tokio::spawn(async move {
      let _data_dir = data_dir;
      server.run(listener).await.expect("Server failed to run");
    });

    Self {
      base_url: format!("http://{}", addr),
      api_secret,
      _handle: Some(handle),
    }
  }

  /// Get the base URL of the test server
  #[allow(unused)]
  pub fn base_url(&self) -> &str {
    &self.base_url
  }

  /// Generate a valid API key for this server
  #[allow(unused)]
  pub fn generate_api_key(&self) -> anyhow::Result<String> {
    let Some(ref secret) = self.api_secret else {
      anyhow::bail!("API secret not configured for this test server");
    };
    Ok(generate_api_key(secret)?)
  }

  /// Create an observation tools client connected to this test server
  #[allow(unused)]
  pub fn create_client(&self) -> anyhow::Result<Client> {
    Ok(ClientBuilder::new().base_url(&self.base_url).build()?)
  }

  /// Create an observation tools client with API key authentication
  #[allow(unused)]
  pub fn create_client_with_api_key(&self, api_key: &str) -> anyhow::Result<Client> {
    Ok(
      ClientBuilder::new()
        .base_url(&self.base_url)
        .api_key(api_key)
        .build()?,
    )
  }

  /// Create an OpenAPI client connected to this test server
  #[allow(unused)]
  pub fn create_api_client(&self) -> anyhow::Result<observation_tools::server_client::Client> {
    observation_tools::server_client::create_client(&self.base_url, None)
  }

  /// Create an OpenAPI client with API key authentication
  #[allow(unused)]
  pub fn create_api_client_with_api_key(
    &self,
    api_key: &str,
  ) -> anyhow::Result<observation_tools::server_client::Client> {
    observation_tools::server_client::create_client(&self.base_url, Some(api_key.to_string()))
  }

  /// List all observations for an execution
  #[allow(unused)]
  pub async fn list_observations(
    &self,
    execution_id: &impl ToString,
  ) -> anyhow::Result<Vec<GetObservation>> {
    let api_client = self.create_api_client()?;
    let response = api_client
      .list_observations()
      .execution_id(&execution_id.to_string())
      .send()
      .await?;
    Ok(response.observations.clone())
  }
}