edgefirst_client/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright © 2025 Au-Zone Technologies. All Rights Reserved.
3
4// SPDX-License-Identifier: Apache-2.0
5// Copyright © 2025 Au-Zone Technologies. All Rights Reserved.
6
7//! # EdgeFirst Studio Client Library
8//!
9//! The EdgeFirst Studio Client Library provides a Rust client for interacting
10//! with EdgeFirst Studio, a comprehensive platform for computer vision and
11//! machine learning workflows. This library enables developers to
12//! programmatically manage datasets, annotations, training sessions, and other
13//! Studio resources.
14//!
15//! ## Features
16//!
17//! - **Authentication**: Secure token-based authentication with automatic
18//!   renewal
19//! - **Dataset Management**: Upload, download, and manage datasets with various
20//!   file types
21//! - **Annotation Management**: Create, update, and retrieve annotations for
22//!   computer vision tasks
23//! - **Training & Validation**: Manage machine learning training and validation
24//!   sessions
25//! - **Project Organization**: Organize work into projects with hierarchical
26//!   structure
27//! - **Polars Integration**: Optional integration with Polars DataFrames for
28//!   data analysis
29//!
30//! ## Quick Start
31//!
32//! ```rust,no_run
33//! use edgefirst_client::{Client, Error};
34//!
35//! #[tokio::main]
36//! async fn main() -> Result<(), Error> {
37//!     // Create a new client
38//!     let client = Client::new()?;
39//!
40//!     // Authenticate with username and password
41//!     let client = client.with_login("username", "password").await?;
42//!
43//!     // List available projects
44//!     let projects = client.projects(None).await?;
45//!     println!("Found {} projects", projects.len());
46//!
47//!     Ok(())
48//! }
49//! ```
50//!
51//! ## Optional Features
52//!
53//! - `polars`: Enables integration with Polars DataFrames for enhanced data
54//!   manipulation
55
56mod api;
57mod client;
58mod dataset;
59mod error;
60mod retry;
61
62pub use crate::{
63    api::{
64        AnnotationSetID, AppId, Artifact, DatasetID, DatasetParams, Experiment, ExperimentID,
65        ImageId, Organization, OrganizationID, Parameter, PresignedUrl, Project, ProjectID,
66        SampleID, SamplesCountResult, SamplesPopulateParams, SamplesPopulateResult, SequenceId,
67        SnapshotID, Stage, Task, TaskID, TaskInfo, TrainingSession, TrainingSessionID,
68        ValidationSession, ValidationSessionID,
69    },
70    client::{Client, Progress},
71    dataset::{
72        Annotation, AnnotationSet, AnnotationType, Box2d, Box3d, Dataset, FileType, GpsData,
73        ImuData, Label, Location, Mask, Sample, SampleFile,
74    },
75    error::Error,
76    retry::{RetryScope, classify_url},
77};
78
79#[cfg(feature = "polars")]
80pub use crate::dataset::annotations_dataframe;
81
82#[cfg(feature = "polars")]
83pub use crate::dataset::samples_dataframe;
84
85#[cfg(feature = "polars")]
86pub use crate::dataset::unflatten_polygon_coordinates;
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use std::{
92        collections::HashMap,
93        env,
94        fs::{File, read_to_string},
95        io::Write,
96        path::PathBuf,
97    };
98
99    /// Get the test data directory (target/testdata)
100    /// Creates it if it doesn't exist
101    fn get_test_data_dir() -> PathBuf {
102        let test_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
103            .parent()
104            .expect("CARGO_MANIFEST_DIR should have parent")
105            .parent()
106            .expect("workspace root should exist")
107            .join("target")
108            .join("testdata");
109
110        std::fs::create_dir_all(&test_dir).expect("Failed to create test data directory");
111        test_dir
112    }
113
114    #[ctor::ctor]
115    fn init() {
116        env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
117    }
118
119    async fn get_client() -> Result<Client, Error> {
120        let client = Client::new()?.with_token_path(None)?;
121
122        let client = match env::var("STUDIO_TOKEN") {
123            Ok(token) => client.with_token(&token)?,
124            Err(_) => client,
125        };
126
127        let client = match env::var("STUDIO_SERVER") {
128            Ok(server) => client.with_server(&server)?,
129            Err(_) => client,
130        };
131
132        let client = match (env::var("STUDIO_USERNAME"), env::var("STUDIO_PASSWORD")) {
133            (Ok(username), Ok(password)) => client.with_login(&username, &password).await?,
134            _ => client,
135        };
136
137        client.verify_token().await?;
138
139        Ok(client)
140    }
141
142    /// Helper: Get training session for "Unit Testing" project
143    async fn get_training_session_for_artifacts() -> Result<TrainingSession, Error> {
144        let client = get_client().await?;
145        let project = client
146            .projects(Some("Unit Testing"))
147            .await?
148            .into_iter()
149            .next()
150            .ok_or_else(|| Error::InvalidParameters("Unit Testing project not found".into()))?;
151        let experiment = client
152            .experiments(project.id(), Some("Unit Testing"))
153            .await?
154            .into_iter()
155            .next()
156            .ok_or_else(|| Error::InvalidParameters("Unit Testing experiment not found".into()))?;
157        let session = client
158            .training_sessions(experiment.id(), Some("modelpack-960x540"))
159            .await?
160            .into_iter()
161            .next()
162            .ok_or_else(|| {
163                Error::InvalidParameters("modelpack-960x540 session not found".into())
164            })?;
165        Ok(session)
166    }
167
168    /// Helper: Get training session for "modelpack-usermanaged"
169    async fn get_training_session_for_checkpoints() -> Result<TrainingSession, Error> {
170        let client = get_client().await?;
171        let project = client
172            .projects(Some("Unit Testing"))
173            .await?
174            .into_iter()
175            .next()
176            .ok_or_else(|| Error::InvalidParameters("Unit Testing project not found".into()))?;
177        let experiment = client
178            .experiments(project.id(), Some("Unit Testing"))
179            .await?
180            .into_iter()
181            .next()
182            .ok_or_else(|| Error::InvalidParameters("Unit Testing experiment not found".into()))?;
183        let session = client
184            .training_sessions(experiment.id(), Some("modelpack-usermanaged"))
185            .await?
186            .into_iter()
187            .next()
188            .ok_or_else(|| {
189                Error::InvalidParameters("modelpack-usermanaged session not found".into())
190            })?;
191        Ok(session)
192    }
193
194    #[tokio::test]
195    async fn test_training_session() -> Result<(), Error> {
196        let client = get_client().await?;
197        let project = client.projects(Some("Unit Testing")).await?;
198        assert!(!project.is_empty());
199        let project = project
200            .first()
201            .expect("'Unit Testing' project should exist");
202        let experiment = client
203            .experiments(project.id(), Some("Unit Testing"))
204            .await?;
205        let experiment = experiment
206            .first()
207            .expect("'Unit Testing' experiment should exist");
208
209        let sessions = client
210            .training_sessions(experiment.id(), Some("modelpack-usermanaged"))
211            .await?;
212        assert_ne!(sessions.len(), 0);
213        let session = sessions
214            .first()
215            .expect("Training sessions should exist for experiment");
216
217        let metrics = HashMap::from([
218            ("epochs".to_string(), Parameter::Integer(10)),
219            ("loss".to_string(), Parameter::Real(0.05)),
220            (
221                "model".to_string(),
222                Parameter::String("modelpack".to_string()),
223            ),
224        ]);
225
226        session.set_metrics(&client, metrics).await?;
227        let updated_metrics = session.metrics(&client).await?;
228        assert_eq!(updated_metrics.len(), 3);
229        assert_eq!(updated_metrics.get("epochs"), Some(&Parameter::Integer(10)));
230        assert_eq!(updated_metrics.get("loss"), Some(&Parameter::Real(0.05)));
231        assert_eq!(
232            updated_metrics.get("model"),
233            Some(&Parameter::String("modelpack".to_string()))
234        );
235
236        println!("Updated Metrics: {:?}", updated_metrics);
237
238        let mut labels = tempfile::NamedTempFile::new()?;
239        write!(labels, "background")?;
240        labels.flush()?;
241
242        session
243            .upload(
244                &client,
245                &[(
246                    "artifacts/labels.txt".to_string(),
247                    labels.path().to_path_buf(),
248                )],
249            )
250            .await?;
251
252        let labels = session.download(&client, "artifacts/labels.txt").await?;
253        assert_eq!(labels, "background");
254
255        Ok(())
256    }
257
258    #[tokio::test]
259    async fn test_validate() -> Result<(), Error> {
260        let client = get_client().await?;
261        let project = client.projects(Some("Unit Testing")).await?;
262        assert!(!project.is_empty());
263        let project = project
264            .first()
265            .expect("'Unit Testing' project should exist");
266
267        let sessions = client.validation_sessions(project.id()).await?;
268        for session in &sessions {
269            let s = client.validation_session(session.id()).await?;
270            assert_eq!(s.id(), session.id());
271            assert_eq!(s.description(), session.description());
272        }
273
274        let session = sessions
275            .into_iter()
276            .find(|s| s.name() == "modelpack-usermanaged")
277            .ok_or_else(|| {
278                Error::InvalidParameters(format!(
279                    "Validation session 'modelpack-usermanaged' not found in project '{}'",
280                    project.name()
281                ))
282            })?;
283
284        let metrics = HashMap::from([("accuracy".to_string(), Parameter::Real(0.95))]);
285        session.set_metrics(&client, metrics).await?;
286
287        let metrics = session.metrics(&client).await?;
288        assert_eq!(metrics.get("accuracy"), Some(&Parameter::Real(0.95)));
289
290        Ok(())
291    }
292
293    #[tokio::test]
294    async fn test_download_artifact_success() -> Result<(), Error> {
295        let trainer = get_training_session_for_artifacts().await?;
296        let client = get_client().await?;
297        let artifacts = client.artifacts(trainer.id()).await?;
298        assert!(!artifacts.is_empty());
299
300        let test_dir = get_test_data_dir();
301        let artifact = &artifacts[0];
302        let output_path = test_dir.join(artifact.name());
303
304        client
305            .download_artifact(
306                trainer.id(),
307                artifact.name(),
308                Some(output_path.clone()),
309                None,
310            )
311            .await?;
312
313        assert!(output_path.exists());
314        if output_path.exists() {
315            std::fs::remove_file(&output_path)?;
316        }
317
318        Ok(())
319    }
320
321    #[tokio::test]
322    async fn test_download_artifact_not_found() -> Result<(), Error> {
323        let trainer = get_training_session_for_artifacts().await?;
324        let client = get_client().await?;
325        let test_dir = get_test_data_dir();
326        let fake_path = test_dir.join("nonexistent_artifact.txt");
327
328        let result = client
329            .download_artifact(
330                trainer.id(),
331                "nonexistent_artifact.txt",
332                Some(fake_path.clone()),
333                None,
334            )
335            .await;
336
337        assert!(result.is_err());
338        assert!(!fake_path.exists());
339
340        Ok(())
341    }
342
343    #[tokio::test]
344    async fn test_artifacts() -> Result<(), Error> {
345        let client = get_client().await?;
346        let project = client.projects(Some("Unit Testing")).await?;
347        assert!(!project.is_empty());
348        let project = project
349            .first()
350            .expect("'Unit Testing' project should exist");
351        let experiment = client
352            .experiments(project.id(), Some("Unit Testing"))
353            .await?;
354        let experiment = experiment
355            .first()
356            .expect("'Unit Testing' experiment should exist");
357        let trainer = client
358            .training_sessions(experiment.id(), Some("modelpack-960x540"))
359            .await?;
360        let trainer = trainer
361            .first()
362            .expect("'modelpack-960x540' training session should exist");
363        let artifacts = client.artifacts(trainer.id()).await?;
364        assert!(!artifacts.is_empty());
365
366        let test_dir = get_test_data_dir();
367
368        for artifact in artifacts {
369            let output_path = test_dir.join(artifact.name());
370            client
371                .download_artifact(
372                    trainer.id(),
373                    artifact.name(),
374                    Some(output_path.clone()),
375                    None,
376                )
377                .await?;
378
379            // Clean up downloaded file
380            if output_path.exists() {
381                std::fs::remove_file(&output_path)?;
382            }
383        }
384
385        let fake_path = test_dir.join("fakefile.txt");
386        let res = client
387            .download_artifact(trainer.id(), "fakefile.txt", Some(fake_path.clone()), None)
388            .await;
389        assert!(res.is_err());
390        assert!(!fake_path.exists());
391
392        Ok(())
393    }
394
395    #[tokio::test]
396    async fn test_download_checkpoint_success() -> Result<(), Error> {
397        let trainer = get_training_session_for_checkpoints().await?;
398        let client = get_client().await?;
399        let test_dir = get_test_data_dir();
400
401        // Create temporary test file
402        let checkpoint_path = test_dir.join("test_checkpoint.txt");
403        {
404            let mut f = File::create(&checkpoint_path)?;
405            f.write_all(b"Test Checkpoint Content")?;
406        }
407
408        // Upload the checkpoint
409        trainer
410            .upload(
411                &client,
412                &[(
413                    "checkpoints/test_checkpoint.txt".to_string(),
414                    checkpoint_path.clone(),
415                )],
416            )
417            .await?;
418
419        // Download and verify
420        let download_path = test_dir.join("downloaded_checkpoint.txt");
421        client
422            .download_checkpoint(
423                trainer.id(),
424                "test_checkpoint.txt",
425                Some(download_path.clone()),
426                None,
427            )
428            .await?;
429
430        let content = read_to_string(&download_path)?;
431        assert_eq!(content, "Test Checkpoint Content");
432
433        // Cleanup
434        if checkpoint_path.exists() {
435            std::fs::remove_file(&checkpoint_path)?;
436        }
437        if download_path.exists() {
438            std::fs::remove_file(&download_path)?;
439        }
440
441        Ok(())
442    }
443
444    #[tokio::test]
445    async fn test_download_checkpoint_not_found() -> Result<(), Error> {
446        let trainer = get_training_session_for_checkpoints().await?;
447        let client = get_client().await?;
448        let test_dir = get_test_data_dir();
449        let fake_path = test_dir.join("nonexistent_checkpoint.txt");
450
451        let result = client
452            .download_checkpoint(
453                trainer.id(),
454                "nonexistent_checkpoint.txt",
455                Some(fake_path.clone()),
456                None,
457            )
458            .await;
459
460        assert!(result.is_err());
461        assert!(!fake_path.exists());
462
463        Ok(())
464    }
465
466    #[tokio::test]
467    async fn test_checkpoints() -> Result<(), Error> {
468        let client = get_client().await?;
469        let project = client.projects(Some("Unit Testing")).await?;
470        assert!(!project.is_empty());
471        let project = project
472            .first()
473            .expect("'Unit Testing' project should exist");
474        let experiment = client
475            .experiments(project.id(), Some("Unit Testing"))
476            .await?;
477        let experiment = experiment.first().ok_or_else(|| {
478            Error::InvalidParameters(format!(
479                "Experiment 'Unit Testing' not found in project '{}'",
480                project.name()
481            ))
482        })?;
483        let trainer = client
484            .training_sessions(experiment.id(), Some("modelpack-usermanaged"))
485            .await?;
486        let trainer = trainer
487            .first()
488            .expect("'modelpack-usermanaged' training session should exist");
489
490        let test_dir = get_test_data_dir();
491        let checkpoint_path = test_dir.join("checkpoint.txt");
492        let checkpoint2_path = test_dir.join("checkpoint2.txt");
493
494        {
495            let mut chkpt = File::create(&checkpoint_path)?;
496            chkpt.write_all(b"Test Checkpoint")?;
497        }
498
499        trainer
500            .upload(
501                &client,
502                &[(
503                    "checkpoints/checkpoint.txt".to_string(),
504                    checkpoint_path.clone(),
505                )],
506            )
507            .await?;
508
509        client
510            .download_checkpoint(
511                trainer.id(),
512                "checkpoint.txt",
513                Some(checkpoint2_path.clone()),
514                None,
515            )
516            .await?;
517
518        let chkpt = read_to_string(&checkpoint2_path)?;
519        assert_eq!(chkpt, "Test Checkpoint");
520
521        let fake_path = test_dir.join("fakefile.txt");
522        let res = client
523            .download_checkpoint(trainer.id(), "fakefile.txt", Some(fake_path.clone()), None)
524            .await;
525        assert!(res.is_err());
526        assert!(!fake_path.exists());
527
528        // Clean up
529        if checkpoint_path.exists() {
530            std::fs::remove_file(&checkpoint_path)?;
531        }
532        if checkpoint2_path.exists() {
533            std::fs::remove_file(&checkpoint2_path)?;
534        }
535
536        Ok(())
537    }
538
539    #[tokio::test]
540    async fn test_task_retrieval() -> Result<(), Error> {
541        let client = get_client().await?;
542
543        // Test: Get all tasks
544        let tasks = client.tasks(None, None, None, None).await?;
545        assert!(!tasks.is_empty());
546
547        // Test: Get task info for first task
548        let task_id = tasks[0].id();
549        let task_info = client.task_info(task_id).await?;
550        assert_eq!(task_info.id(), task_id);
551
552        Ok(())
553    }
554
555    #[tokio::test]
556    async fn test_task_filtering_by_name() -> Result<(), Error> {
557        let client = get_client().await?;
558        let project = client.projects(Some("Unit Testing")).await?;
559        let project = project
560            .first()
561            .expect("'Unit Testing' project should exist");
562
563        // Test: Get tasks by name
564        let tasks = client
565            .tasks(Some("modelpack-usermanaged"), None, None, None)
566            .await?;
567
568        if !tasks.is_empty() {
569            // Get detailed info for each task
570            let task_infos = tasks
571                .into_iter()
572                .map(|t| client.task_info(t.id()))
573                .collect::<Vec<_>>();
574            let task_infos = futures::future::try_join_all(task_infos).await?;
575
576            // Filter by project
577            let filtered = task_infos
578                .into_iter()
579                .filter(|t| t.project_id() == Some(project.id()))
580                .collect::<Vec<_>>();
581
582            if !filtered.is_empty() {
583                assert_eq!(filtered[0].project_id(), Some(project.id()));
584            }
585        }
586
587        Ok(())
588    }
589
590    #[tokio::test]
591    async fn test_task_status_and_stages() -> Result<(), Error> {
592        let client = get_client().await?;
593
594        // Get first available task
595        let tasks = client.tasks(None, None, None, None).await?;
596        if tasks.is_empty() {
597            return Ok(());
598        }
599
600        let task_id = tasks[0].id();
601
602        // Test: Get task status
603        let status = client.task_status(task_id, "training").await?;
604        assert_eq!(status.id(), task_id);
605        assert_eq!(status.status(), "training");
606
607        // Test: Set stages
608        let stages = [
609            ("download", "Downloading Dataset"),
610            ("train", "Training Model"),
611            ("export", "Exporting Model"),
612        ];
613        client.set_stages(task_id, &stages).await?;
614
615        // Test: Update stage
616        client
617            .update_stage(task_id, "download", "running", "Downloading dataset", 50)
618            .await?;
619
620        // Verify task with updated stages
621        let updated_task = client.task_info(task_id).await?;
622        assert_eq!(updated_task.id(), task_id);
623
624        Ok(())
625    }
626
627    #[tokio::test]
628    async fn test_tasks() -> Result<(), Error> {
629        let client = get_client().await?;
630        let project = client.projects(Some("Unit Testing")).await?;
631        let project = project
632            .first()
633            .expect("'Unit Testing' project should exist");
634        let tasks = client.tasks(None, None, None, None).await?;
635
636        for task in tasks {
637            let task_info = client.task_info(task.id()).await?;
638            println!("{} - {}", task, task_info);
639        }
640
641        let tasks = client
642            .tasks(Some("modelpack-usermanaged"), None, None, None)
643            .await?;
644        let tasks = tasks
645            .into_iter()
646            .map(|t| client.task_info(t.id()))
647            .collect::<Vec<_>>();
648        let tasks = futures::future::try_join_all(tasks).await?;
649        let tasks = tasks
650            .into_iter()
651            .filter(|t| t.project_id() == Some(project.id()))
652            .collect::<Vec<_>>();
653        assert_ne!(tasks.len(), 0);
654        let task = &tasks[0];
655
656        let t = client.task_status(task.id(), "training").await?;
657        assert_eq!(t.id(), task.id());
658        assert_eq!(t.status(), "training");
659
660        let stages = [
661            ("download", "Downloading Dataset"),
662            ("train", "Training Model"),
663            ("export", "Exporting Model"),
664        ];
665        client.set_stages(task.id(), &stages).await?;
666
667        client
668            .update_stage(task.id(), "download", "running", "Downloading dataset", 50)
669            .await?;
670
671        let task = client.task_info(task.id()).await?;
672        println!("task progress: {:?}", task.stages());
673
674        Ok(())
675    }
676
677    // ============================================================================
678    // Retry URL Classification Tests
679    // ============================================================================
680
681    mod retry_url_classification {
682        use super::*;
683
684        #[test]
685        fn test_studio_api_base_url() {
686            // Base production URL
687            assert_eq!(
688                classify_url("https://edgefirst.studio/api"),
689                RetryScope::StudioApi
690            );
691        }
692
693        #[test]
694        fn test_studio_api_with_trailing_slash() {
695            // Trailing slash should be handled correctly
696            assert_eq!(
697                classify_url("https://edgefirst.studio/api/"),
698                RetryScope::StudioApi
699            );
700        }
701
702        #[test]
703        fn test_studio_api_with_path() {
704            // API endpoints with additional path segments
705            assert_eq!(
706                classify_url("https://edgefirst.studio/api/datasets"),
707                RetryScope::StudioApi
708            );
709            assert_eq!(
710                classify_url("https://edgefirst.studio/api/auth.login"),
711                RetryScope::StudioApi
712            );
713            assert_eq!(
714                classify_url("https://edgefirst.studio/api/trainer/session"),
715                RetryScope::StudioApi
716            );
717        }
718
719        #[test]
720        fn test_studio_api_with_query_params() {
721            // Query parameters should not affect classification
722            assert_eq!(
723                classify_url("https://edgefirst.studio/api?foo=bar"),
724                RetryScope::StudioApi
725            );
726            assert_eq!(
727                classify_url("https://edgefirst.studio/api/datasets?page=1&limit=10"),
728                RetryScope::StudioApi
729            );
730        }
731
732        #[test]
733        fn test_studio_api_subdomains() {
734            // Server-specific instances (test, stage, saas, ocean, etc.)
735            assert_eq!(
736                classify_url("https://test.edgefirst.studio/api"),
737                RetryScope::StudioApi
738            );
739            assert_eq!(
740                classify_url("https://stage.edgefirst.studio/api"),
741                RetryScope::StudioApi
742            );
743            assert_eq!(
744                classify_url("https://saas.edgefirst.studio/api"),
745                RetryScope::StudioApi
746            );
747            assert_eq!(
748                classify_url("https://ocean.edgefirst.studio/api"),
749                RetryScope::StudioApi
750            );
751        }
752
753        #[test]
754        fn test_studio_api_with_standard_port() {
755            // Standard HTTPS port (443) should be handled
756            assert_eq!(
757                classify_url("https://edgefirst.studio:443/api"),
758                RetryScope::StudioApi
759            );
760            assert_eq!(
761                classify_url("https://test.edgefirst.studio:443/api"),
762                RetryScope::StudioApi
763            );
764        }
765
766        #[test]
767        fn test_studio_api_with_custom_port() {
768            // Custom ports should be handled correctly
769            assert_eq!(
770                classify_url("https://test.edgefirst.studio:8080/api"),
771                RetryScope::StudioApi
772            );
773            assert_eq!(
774                classify_url("https://edgefirst.studio:8443/api"),
775                RetryScope::StudioApi
776            );
777        }
778
779        #[test]
780        fn test_studio_api_http_protocol() {
781            // HTTP (not HTTPS) should still be recognized
782            assert_eq!(
783                classify_url("http://edgefirst.studio/api"),
784                RetryScope::StudioApi
785            );
786            assert_eq!(
787                classify_url("http://test.edgefirst.studio/api"),
788                RetryScope::StudioApi
789            );
790        }
791
792        #[test]
793        fn test_file_io_s3_urls() {
794            // S3 URLs for file operations
795            assert_eq!(
796                classify_url("https://s3.amazonaws.com/bucket/file.bin"),
797                RetryScope::FileIO
798            );
799            assert_eq!(
800                classify_url("https://s3.us-west-2.amazonaws.com/mybucket/data.zip"),
801                RetryScope::FileIO
802            );
803        }
804
805        #[test]
806        fn test_file_io_cloudfront_urls() {
807            // CloudFront URLs for file distribution
808            assert_eq!(
809                classify_url("https://d123abc.cloudfront.net/file.bin"),
810                RetryScope::FileIO
811            );
812            assert_eq!(
813                classify_url("https://d456def.cloudfront.net/path/to/file.tar.gz"),
814                RetryScope::FileIO
815            );
816        }
817
818        #[test]
819        fn test_file_io_non_api_studio_paths() {
820            // Non-API paths on edgefirst.studio domain
821            assert_eq!(
822                classify_url("https://edgefirst.studio/docs"),
823                RetryScope::FileIO
824            );
825            assert_eq!(
826                classify_url("https://edgefirst.studio/download_model"),
827                RetryScope::FileIO
828            );
829            assert_eq!(
830                classify_url("https://test.edgefirst.studio/download_model"),
831                RetryScope::FileIO
832            );
833            assert_eq!(
834                classify_url("https://stage.edgefirst.studio/download_checkpoint"),
835                RetryScope::FileIO
836            );
837        }
838
839        #[test]
840        fn test_file_io_generic_urls() {
841            // Generic download URLs
842            assert_eq!(
843                classify_url("https://example.com/download"),
844                RetryScope::FileIO
845            );
846            assert_eq!(
847                classify_url("https://cdn.example.com/files/data.json"),
848                RetryScope::FileIO
849            );
850        }
851
852        #[test]
853        fn test_security_malicious_url_substring() {
854            // Security: URL with edgefirst.studio in path should NOT match
855            assert_eq!(
856                classify_url("https://evil.com/test.edgefirst.studio/api"),
857                RetryScope::FileIO
858            );
859            assert_eq!(
860                classify_url("https://attacker.com/edgefirst.studio/api/fake"),
861                RetryScope::FileIO
862            );
863        }
864
865        #[test]
866        fn test_edge_case_similar_domains() {
867            // Similar but different domains should be FileIO
868            assert_eq!(
869                classify_url("https://edgefirst.studio.com/api"),
870                RetryScope::FileIO
871            );
872            assert_eq!(
873                classify_url("https://notedgefirst.studio/api"),
874                RetryScope::FileIO
875            );
876            assert_eq!(
877                classify_url("https://edgefirststudio.com/api"),
878                RetryScope::FileIO
879            );
880        }
881
882        #[test]
883        fn test_edge_case_invalid_urls() {
884            // Invalid URLs should default to FileIO
885            assert_eq!(classify_url("not a url"), RetryScope::FileIO);
886            assert_eq!(classify_url(""), RetryScope::FileIO);
887            assert_eq!(
888                classify_url("ftp://edgefirst.studio/api"),
889                RetryScope::FileIO
890            );
891        }
892
893        #[test]
894        fn test_edge_case_url_normalization() {
895            // URL normalization edge cases
896            assert_eq!(
897                classify_url("https://EDGEFIRST.STUDIO/api"),
898                RetryScope::StudioApi
899            );
900            assert_eq!(
901                classify_url("https://test.EDGEFIRST.studio/api"),
902                RetryScope::StudioApi
903            );
904        }
905
906        #[test]
907        fn test_comprehensive_subdomain_coverage() {
908            // Ensure all known server instances are recognized
909            let subdomains = vec![
910                "test", "stage", "saas", "ocean", "prod", "dev", "qa", "demo",
911            ];
912
913            for subdomain in subdomains {
914                let url = format!("https://{}.edgefirst.studio/api", subdomain);
915                assert_eq!(
916                    classify_url(&url),
917                    RetryScope::StudioApi,
918                    "Failed for subdomain: {}",
919                    subdomain
920                );
921            }
922        }
923
924        #[test]
925        fn test_api_path_variations() {
926            // Various API path patterns
927            assert_eq!(
928                classify_url("https://edgefirst.studio/api"),
929                RetryScope::StudioApi
930            );
931            assert_eq!(
932                classify_url("https://edgefirst.studio/api/"),
933                RetryScope::StudioApi
934            );
935            assert_eq!(
936                classify_url("https://edgefirst.studio/api/v1"),
937                RetryScope::StudioApi
938            );
939            assert_eq!(
940                classify_url("https://edgefirst.studio/api/v2/datasets"),
941                RetryScope::StudioApi
942            );
943
944            // Non-/api paths should be FileIO
945            assert_eq!(
946                classify_url("https://edgefirst.studio/apis"),
947                RetryScope::FileIO
948            );
949            assert_eq!(
950                classify_url("https://edgefirst.studio/v1/api"),
951                RetryScope::FileIO
952            );
953        }
954
955        #[test]
956        fn test_port_range_coverage() {
957            // Test various port numbers
958            let ports = vec![80, 443, 8080, 8443, 3000, 5000, 9000];
959
960            for port in ports {
961                let url = format!("https://test.edgefirst.studio:{}/api", port);
962                assert_eq!(
963                    classify_url(&url),
964                    RetryScope::StudioApi,
965                    "Failed for port: {}",
966                    port
967                );
968            }
969        }
970
971        #[test]
972        fn test_complex_query_strings() {
973            // Complex query parameters with special characters
974            assert_eq!(
975                classify_url("https://edgefirst.studio/api?token=abc123&redirect=/dashboard"),
976                RetryScope::StudioApi
977            );
978            assert_eq!(
979                classify_url("https://test.edgefirst.studio/api?q=search%20term&page=1"),
980                RetryScope::StudioApi
981            );
982        }
983
984        #[test]
985        fn test_url_with_fragment() {
986            // URLs with fragments (#) - fragments are not sent to server
987            assert_eq!(
988                classify_url("https://edgefirst.studio/api#section"),
989                RetryScope::StudioApi
990            );
991            assert_eq!(
992                classify_url("https://test.edgefirst.studio/api/datasets#results"),
993                RetryScope::StudioApi
994            );
995        }
996    }
997}