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        Snapshot, SnapshotID, SnapshotRestoreResult, Stage, Task, TaskID, TaskInfo,
68        TrainingSession, TrainingSessionID, 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 tasks = client.tasks(None, None, None, None).await?;
631
632        for task in tasks {
633            let task_info = client.task_info(task.id()).await?;
634            println!("{} - {}", task, task_info);
635        }
636
637        let tasks = client
638            .tasks(Some("modelpack-usermanaged"), None, None, None)
639            .await?;
640        let tasks = tasks
641            .into_iter()
642            .map(|t| client.task_info(t.id()))
643            .collect::<Vec<_>>();
644        let tasks = futures::future::try_join_all(tasks).await?;
645        assert_ne!(tasks.len(), 0);
646        let task = &tasks[0];
647
648        let t = client.task_status(task.id(), "training").await?;
649        assert_eq!(t.id(), task.id());
650        assert_eq!(t.status(), "training");
651
652        let stages = [
653            ("download", "Downloading Dataset"),
654            ("train", "Training Model"),
655            ("export", "Exporting Model"),
656        ];
657        client.set_stages(task.id(), &stages).await?;
658
659        client
660            .update_stage(task.id(), "download", "running", "Downloading dataset", 50)
661            .await?;
662
663        let task = client.task_info(task.id()).await?;
664        println!("task progress: {:?}", task.stages());
665
666        Ok(())
667    }
668
669    // ============================================================================
670    // Retry URL Classification Tests
671    // ============================================================================
672
673    mod retry_url_classification {
674        use super::*;
675
676        #[test]
677        fn test_studio_api_base_url() {
678            // Base production URL
679            assert_eq!(
680                classify_url("https://edgefirst.studio/api"),
681                RetryScope::StudioApi
682            );
683        }
684
685        #[test]
686        fn test_studio_api_with_trailing_slash() {
687            // Trailing slash should be handled correctly
688            assert_eq!(
689                classify_url("https://edgefirst.studio/api/"),
690                RetryScope::StudioApi
691            );
692        }
693
694        #[test]
695        fn test_studio_api_with_path() {
696            // API endpoints with additional path segments
697            assert_eq!(
698                classify_url("https://edgefirst.studio/api/datasets"),
699                RetryScope::StudioApi
700            );
701            assert_eq!(
702                classify_url("https://edgefirst.studio/api/auth.login"),
703                RetryScope::StudioApi
704            );
705            assert_eq!(
706                classify_url("https://edgefirst.studio/api/trainer/session"),
707                RetryScope::StudioApi
708            );
709        }
710
711        #[test]
712        fn test_studio_api_with_query_params() {
713            // Query parameters should not affect classification
714            assert_eq!(
715                classify_url("https://edgefirst.studio/api?foo=bar"),
716                RetryScope::StudioApi
717            );
718            assert_eq!(
719                classify_url("https://edgefirst.studio/api/datasets?page=1&limit=10"),
720                RetryScope::StudioApi
721            );
722        }
723
724        #[test]
725        fn test_studio_api_subdomains() {
726            // Server-specific instances (test, stage, saas, ocean, etc.)
727            assert_eq!(
728                classify_url("https://test.edgefirst.studio/api"),
729                RetryScope::StudioApi
730            );
731            assert_eq!(
732                classify_url("https://stage.edgefirst.studio/api"),
733                RetryScope::StudioApi
734            );
735            assert_eq!(
736                classify_url("https://saas.edgefirst.studio/api"),
737                RetryScope::StudioApi
738            );
739            assert_eq!(
740                classify_url("https://ocean.edgefirst.studio/api"),
741                RetryScope::StudioApi
742            );
743        }
744
745        #[test]
746        fn test_studio_api_with_standard_port() {
747            // Standard HTTPS port (443) should be handled
748            assert_eq!(
749                classify_url("https://edgefirst.studio:443/api"),
750                RetryScope::StudioApi
751            );
752            assert_eq!(
753                classify_url("https://test.edgefirst.studio:443/api"),
754                RetryScope::StudioApi
755            );
756        }
757
758        #[test]
759        fn test_studio_api_with_custom_port() {
760            // Custom ports should be handled correctly
761            assert_eq!(
762                classify_url("https://test.edgefirst.studio:8080/api"),
763                RetryScope::StudioApi
764            );
765            assert_eq!(
766                classify_url("https://edgefirst.studio:8443/api"),
767                RetryScope::StudioApi
768            );
769        }
770
771        #[test]
772        fn test_studio_api_http_protocol() {
773            // HTTP (not HTTPS) should still be recognized
774            assert_eq!(
775                classify_url("http://edgefirst.studio/api"),
776                RetryScope::StudioApi
777            );
778            assert_eq!(
779                classify_url("http://test.edgefirst.studio/api"),
780                RetryScope::StudioApi
781            );
782        }
783
784        #[test]
785        fn test_file_io_s3_urls() {
786            // S3 URLs for file operations
787            assert_eq!(
788                classify_url("https://s3.amazonaws.com/bucket/file.bin"),
789                RetryScope::FileIO
790            );
791            assert_eq!(
792                classify_url("https://s3.us-west-2.amazonaws.com/mybucket/data.zip"),
793                RetryScope::FileIO
794            );
795        }
796
797        #[test]
798        fn test_file_io_cloudfront_urls() {
799            // CloudFront URLs for file distribution
800            assert_eq!(
801                classify_url("https://d123abc.cloudfront.net/file.bin"),
802                RetryScope::FileIO
803            );
804            assert_eq!(
805                classify_url("https://d456def.cloudfront.net/path/to/file.tar.gz"),
806                RetryScope::FileIO
807            );
808        }
809
810        #[test]
811        fn test_file_io_non_api_studio_paths() {
812            // Non-API paths on edgefirst.studio domain
813            assert_eq!(
814                classify_url("https://edgefirst.studio/docs"),
815                RetryScope::FileIO
816            );
817            assert_eq!(
818                classify_url("https://edgefirst.studio/download_model"),
819                RetryScope::FileIO
820            );
821            assert_eq!(
822                classify_url("https://test.edgefirst.studio/download_model"),
823                RetryScope::FileIO
824            );
825            assert_eq!(
826                classify_url("https://stage.edgefirst.studio/download_checkpoint"),
827                RetryScope::FileIO
828            );
829        }
830
831        #[test]
832        fn test_file_io_generic_urls() {
833            // Generic download URLs
834            assert_eq!(
835                classify_url("https://example.com/download"),
836                RetryScope::FileIO
837            );
838            assert_eq!(
839                classify_url("https://cdn.example.com/files/data.json"),
840                RetryScope::FileIO
841            );
842        }
843
844        #[test]
845        fn test_security_malicious_url_substring() {
846            // Security: URL with edgefirst.studio in path should NOT match
847            assert_eq!(
848                classify_url("https://evil.com/test.edgefirst.studio/api"),
849                RetryScope::FileIO
850            );
851            assert_eq!(
852                classify_url("https://attacker.com/edgefirst.studio/api/fake"),
853                RetryScope::FileIO
854            );
855        }
856
857        #[test]
858        fn test_edge_case_similar_domains() {
859            // Similar but different domains should be FileIO
860            assert_eq!(
861                classify_url("https://edgefirst.studio.com/api"),
862                RetryScope::FileIO
863            );
864            assert_eq!(
865                classify_url("https://notedgefirst.studio/api"),
866                RetryScope::FileIO
867            );
868            assert_eq!(
869                classify_url("https://edgefirststudio.com/api"),
870                RetryScope::FileIO
871            );
872        }
873
874        #[test]
875        fn test_edge_case_invalid_urls() {
876            // Invalid URLs should default to FileIO
877            assert_eq!(classify_url("not a url"), RetryScope::FileIO);
878            assert_eq!(classify_url(""), RetryScope::FileIO);
879            assert_eq!(
880                classify_url("ftp://edgefirst.studio/api"),
881                RetryScope::FileIO
882            );
883        }
884
885        #[test]
886        fn test_edge_case_url_normalization() {
887            // URL normalization edge cases
888            assert_eq!(
889                classify_url("https://EDGEFIRST.STUDIO/api"),
890                RetryScope::StudioApi
891            );
892            assert_eq!(
893                classify_url("https://test.EDGEFIRST.studio/api"),
894                RetryScope::StudioApi
895            );
896        }
897
898        #[test]
899        fn test_comprehensive_subdomain_coverage() {
900            // Ensure all known server instances are recognized
901            let subdomains = vec![
902                "test", "stage", "saas", "ocean", "prod", "dev", "qa", "demo",
903            ];
904
905            for subdomain in subdomains {
906                let url = format!("https://{}.edgefirst.studio/api", subdomain);
907                assert_eq!(
908                    classify_url(&url),
909                    RetryScope::StudioApi,
910                    "Failed for subdomain: {}",
911                    subdomain
912                );
913            }
914        }
915
916        #[test]
917        fn test_api_path_variations() {
918            // Various API path patterns
919            assert_eq!(
920                classify_url("https://edgefirst.studio/api"),
921                RetryScope::StudioApi
922            );
923            assert_eq!(
924                classify_url("https://edgefirst.studio/api/"),
925                RetryScope::StudioApi
926            );
927            assert_eq!(
928                classify_url("https://edgefirst.studio/api/v1"),
929                RetryScope::StudioApi
930            );
931            assert_eq!(
932                classify_url("https://edgefirst.studio/api/v2/datasets"),
933                RetryScope::StudioApi
934            );
935
936            // Non-/api paths should be FileIO
937            assert_eq!(
938                classify_url("https://edgefirst.studio/apis"),
939                RetryScope::FileIO
940            );
941            assert_eq!(
942                classify_url("https://edgefirst.studio/v1/api"),
943                RetryScope::FileIO
944            );
945        }
946
947        #[test]
948        fn test_port_range_coverage() {
949            // Test various port numbers
950            let ports = vec![80, 443, 8080, 8443, 3000, 5000, 9000];
951
952            for port in ports {
953                let url = format!("https://test.edgefirst.studio:{}/api", port);
954                assert_eq!(
955                    classify_url(&url),
956                    RetryScope::StudioApi,
957                    "Failed for port: {}",
958                    port
959                );
960            }
961        }
962
963        #[test]
964        fn test_complex_query_strings() {
965            // Complex query parameters with special characters
966            assert_eq!(
967                classify_url("https://edgefirst.studio/api?token=abc123&redirect=/dashboard"),
968                RetryScope::StudioApi
969            );
970            assert_eq!(
971                classify_url("https://test.edgefirst.studio/api?q=search%20term&page=1"),
972                RetryScope::StudioApi
973            );
974        }
975
976        #[test]
977        fn test_url_with_fragment() {
978            // URLs with fragments (#) - fragments are not sent to server
979            assert_eq!(
980                classify_url("https://edgefirst.studio/api#section"),
981                RetryScope::StudioApi
982            );
983            assert_eq!(
984                classify_url("https://test.edgefirst.studio/api/datasets#results"),
985                RetryScope::StudioApi
986            );
987        }
988    }
989}