antimatter 2.0.13

use antimatter::capsule::common::{CellReader, Column, RowReader, SpanTag};
use antimatter::session::api_helper::domains;
use antimatter::session::session::SessionError::APIError;
use antimatter::session::session::{EncapsulateConfig, SessionError};
use antimatter_api::models;
use antimatter_api::models::data_policy_clause::Operator;
use antimatter_api::models::fact_policy_rules_inner_arguments_inner::Source::Any;
use antimatter_api::models::new_data_policy_rule::{TokenFormat, TokenScope};
use antimatter_api::models::root_encryption_key_test_response::Status::Healthy;
use antimatter_api::models::AddReadContext;
use antimatter_api::models::{
    fact_policy_rules_inner, FactPolicyRulesInner, FactPolicyRulesInnerArgumentsInner,
    NewDomainPolicyRule, PolicyRuleOperation, PolicyRuleResult,
};
use antimatter_api::models::{ActiveRootEncryptionKeyId, FactTuple};
use antimatter_api::models::{
    AddWriteContext, CapabilityExpression, ClassifierRule, DataPolicyClause, DataPolicyRuleChanges,
    DataPolicyRuleEffect, FactExpression, FactExpressionArgumentsInner, LlmClassifierConfig,
    NewDataPolicy, NewDataPolicyRule, RegexClassifierConfig, SetDataPolicyBinding,
    SetDataPolicyBindingReadContextsInner, TagExpression, WriteContextConfigInfo,
};
use antimatter_api::models::{
    Capability, DomainIdentityEmailPrincipalParams, DomainIdentityPrincipalDetails,
    DomainIdentityProviderDetails, DomainIdentityProviderPrincipalParams,
    GoogleOAuthDomainIdentityProviderDetails,
    GoogleOAuthDomainIdentityProviderDetailsGroupMappings,
    GoogleOAuthDomainIdentityProviderGroupCapabilityMappings,
    GoogleOAuthDomainIdentityProviderGroupMappingDetails, UpdatePrincipalParams,
};
use antimatter_api::models::{DeleteTags, TagTypeField};
use antimatter_api::models::{NewFact, NewFactTypeDefinition, NewFactTypeDefinitionArgumentsInner};
use std::collections::HashMap;
use std::io::Cursor;

#[test]
fn test_api_write_contexts() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    let res = session
        .list_write_context()
        .expect("failed to list write contexts");
    assert_eq!(
        res.write_contexts[0].name, "default",
        "unexpected write context name"
    );

    let res = session
        .describe_write_context("default")
        .expect("failed to describe write context");
    assert_eq!(res.imported, false, "unexpected imported value");

    session
        .add_write_context(
            "test_write_context1",
            AddWriteContext {
                summary: "summary1".to_string(),
                description: "description1".to_string(),
                config: Box::new(WriteContextConfigInfo {
                    key_reuse_ttl: Some(0),
                    default_capsule_tags: Some(vec![]),
                    required_hooks: vec![],
                }),
            },
        )
        .expect("failed to add write context");

    let res = session
        .list_write_context()
        .expect("failed to list write contexts");
    assert_eq!(
        res.write_contexts.len(),
        3,
        "unexpected write context count"
    );

    session
        .upsert_write_context_configuration(
            "test_write_context1",
            WriteContextConfigInfo {
                key_reuse_ttl: Some(60),
                default_capsule_tags: Some(vec![]),
                required_hooks: vec![],
            },
        )
        .expect("failed to upsert write context");

    let res = session
        .describe_write_context("test_write_context1")
        .expect("failed to describe write context");
    assert_eq!(
        res.config.key_reuse_ttl.unwrap_or(0),
        60,
        "unexpected imported value"
    );

    session
        .insert_write_context_classifier_rule(
            "test_write_context1",
            ClassifierRule {
                id: None,
                comment: None,
                span_tags: vec![],
                capsule_tags: vec![],
                regex_config: Some(Box::new(RegexClassifierConfig {
                    pattern: "[0-9]{{3}}-[0-9]{{3}}-[0-9]{{4}}".to_string(),
                    match_on_key: false,
                })),
                llm_config: None,
            },
        )
        .expect("failed to insert write context classifier rule");

    session
        .insert_write_context_classifier_rule(
            "test_write_context1",
            ClassifierRule {
                id: None,
                comment: None,
                span_tags: vec![],
                capsule_tags: vec![],
                regex_config: None,
                llm_config: Some(Box::new(LlmClassifierConfig {
                    model: "gpt-4".to_string(),
                    prompt: "this is a test prompt".to_string(),
                })),
            },
        )
        .expect("failed to insert write context classifier rule");

    session
        .insert_write_context_classifier_rule(
            "test_write_context1",
            ClassifierRule {
                id: None,
                comment: None,
                span_tags: vec![],
                capsule_tags: vec![],
                regex_config: Some(Box::new(RegexClassifierConfig {
                    pattern: "[0-9]{{3}}".to_string(),
                    match_on_key: false,
                })),
                llm_config: None,
            },
        )
        .expect("failed to insert write context classifier rule");

    let res = session
        .list_write_context_classifier_rules("test_write_context1")
        .expect("failed to list write context classifier rules");
    let x = res.rules.unwrap_or_default();
    assert_eq!(x.len(), 3, "expected 3 write context classifier rules");

    session
        .delete_write_context_classifier_rule(
            "test_write_context1",
            x[0].id.clone().unwrap_or_default().as_str(),
        )
        .expect("failed to delete write context classifier rule");

    let res = session
        .list_write_context_classifier_rules("test_write_context1")
        .expect("failed to list write context classifier rules");
    assert_eq!(
        res.rules.unwrap_or_default().len(),
        2,
        "expected 2 write context classifier rules"
    );

    session
        .delete_write_context_classifier_rules("test_write_context1")
        .expect("failed to delete write context classifier rules");

    let res = session
        .list_write_context_classifier_rules("test_write_context1")
        .expect("failed to list write context classifier rules");
    assert_eq!(
        res.rules.unwrap_or_default().len(),
        0,
        "expected 0 write context classifier rules"
    );

    session
        .delete_write_context("test_write_context1")
        .expect("failed to delete write context");

    let res = session
        .list_write_context()
        .expect("failed to list write contexts");
    assert_eq!(
        res.write_contexts.len(),
        2,
        "unexpected write context count"
    );
}

#[test]
#[ignore]
fn test_api_verification() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    // NOTE: The 400 is expected here as the domain has no emails pending.
    let _ = session
        .resend_verification_email("test@antimatter.io")
        .map_err(|e| {
            assert_eq!(
                e,
                APIError("error in response: status code 400 Bad Request".to_string())
            )
        });
}

#[test]
fn test_api_root_keys() {
    // TODO: Add in a test for `add_root_encryption_key`, this will require
    //  either setting setting the GCP key's label, or AWS key's alias. The
    //  same should be done for `delete_root_encryption_key`. For now it is okay
    //  though as both v1 and v2 API integration do a full test of this
    //  functionality.

    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    let res = session
        .list_root_encryption_keys()
        .expect("failed to fetch REK list");
    assert_eq!(
        res.keys.unwrap_or_default().len(),
        1,
        "unexpected REK count"
    );

    let res = session
        .get_active_root_encryption_key()
        .expect("failed to fetch active REK");
    assert_eq!(res.source, "default", "unexpected REK source");

    let res = session
        .test_root_encryption_key(res.rek_id.as_str())
        .expect("failed to test REK");
    assert_eq!(res.status, Healthy, "expected REK to be HEALTHY");

    session
        .set_active_root_encryption_key(ActiveRootEncryptionKeyId {
            key_id: res.id.to_string(),
            rotate_batch: None,
        })
        .expect("failed to set active REK");

    let res = session
        .rotate_encryption_keys()
        .expect("failed to rotate KEKs");
    assert_eq!(res.has_more, false, "unexpected has_more value");

    let res = session
        .list_key_providers()
        .expect("failed to list key providers");
    assert_eq!(
        res.providers.unwrap_or_default().len() >= 2,
        true,
        "expected at least 2 key providers"
    );
}

#[test]
fn test_api_data_policy() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    session
        .add_read_context(
            "testreadctx1",
            AddReadContext {
                summary: "read context summary1".to_string(),
                description: "read context description1".to_string(),
                disable_read_logging: Some(false),
                key_cache_ttl: Some(0),
                required_hooks: Some(vec![]),
                read_parameters: Some(vec![]),
            },
        )
        .expect("failed to add read context");

    let res = session
        .describe_read_context("testreadctx1", Some(false))
        .expect("failed to get read context");
    assert_eq!(
        res.summary, "read context summary1",
        "unexpected read context summary"
    );

    let policy_id = session
        .create_data_policy(NewDataPolicy {
            name: "testpolicy1".to_string(),
            description: "test policy".to_string(),
        })
        .expect("failed to create data policy")
        .policy_id;

    let rule_1_id = session
        .update_data_policy_rules(
            policy_id.as_str(),
            DataPolicyRuleChanges {
                delete_rules: None,
                new_rules: Some(vec![NewDataPolicyRule {
                    comment: None,
                    clauses: vec![DataPolicyClause {
                        operator: Operator::AllOf,
                        tags: None,
                        capabilities: None,
                        facts: Some(vec![
                            FactExpression{
                                r#type: "fact_type".to_string(),
                                operator: antimatter_api::models::fact_expression::Operator::Exists,
                                arguments: vec![FactExpressionArgumentsInner{
                                    operator: antimatter_api::models::fact_expression_arguments_inner::Operator::In,
                                    values: Some(vec!["test".to_string()]),
                                }],
                                variables: None,
                            },
                        ]),
                        read_parameters: None,
                    }],
                    effect: DataPolicyRuleEffect::Allow,
                    token_scope: Some(TokenScope::Capsule),
                    token_format: Some(TokenFormat::Synthetic),
                    priority: Some(0),
                    assign_priority: None,
                }]),
            },
        )
        .expect("failed to add data policy rule #1").new_rules[0].clone();

    let rule_2_id = session
        .update_data_policy_rules(
            policy_id.as_str(),
            DataPolicyRuleChanges {
                delete_rules: None,
                new_rules: Some(vec![NewDataPolicyRule {
                    comment: None,
                    clauses: vec![DataPolicyClause {
                        operator: Operator::AllOf,
                        tags: None,
                        capabilities: Some(vec![CapabilityExpression {
                            name: "test_cap".to_string(),
                            values: Some(vec!["test_value".to_string()]),
                            operator: antimatter_api::models::capability_expression::Operator::In,
                            variables: None,
                        }]),
                        facts: None,
                        read_parameters: None,
                    }],
                    effect: DataPolicyRuleEffect::Allow,
                    token_scope: Some(TokenScope::Capsule),
                    token_format: Some(TokenFormat::Synthetic),
                    priority: Some(0),
                    assign_priority: None,
                }]),
            },
        )
        .expect("failed to add data policy rule #2")
        .new_rules[0]
        .clone();

    let rule_3_id = session
        .update_data_policy_rules(
            policy_id.as_str(),
            DataPolicyRuleChanges {
                delete_rules: None,
                new_rules: Some(vec![NewDataPolicyRule {
                    comment: None,
                    clauses: vec![DataPolicyClause {
                        operator: Operator::AllOf,
                        tags: Some(vec![TagExpression {
                            name: "antimatter.io/test".to_string(),
                            values: Some(vec!["test_value".to_string()]),
                            operator: antimatter_api::models::tag_expression::Operator::In,
                            variables: None,
                        }]),
                        capabilities: None,
                        facts: None,
                        read_parameters: None,
                    }],
                    effect: DataPolicyRuleEffect::Allow,
                    token_scope: Some(TokenScope::Capsule),
                    token_format: Some(TokenFormat::Synthetic),
                    priority: Some(0),
                    assign_priority: None,
                }]),
            },
        )
        .expect("failed to add data policy rule #3")
        .new_rules[0]
        .clone();

    session
        .set_data_policy_binding(
            policy_id.as_str(),
            SetDataPolicyBinding {
                read_contexts: Some(vec![SetDataPolicyBindingReadContextsInner {
                    name: "testreadctx1".to_string(),
                    configuration: antimatter_api::models::set_data_policy_binding_read_contexts_inner::Configuration::Attached,
                }]),
                default_attachment:
                antimatter_api::models::set_data_policy_binding::DefaultAttachment::Attached,
            },
        )
        .expect("failed to set data policy binding");

    session
        .describe_data_policy_rule(policy_id.as_str(), rule_1_id.as_str())
        .expect("failed to get rule 1");
    session
        .describe_data_policy_rule(policy_id.as_str(), rule_2_id.as_str())
        .expect("failed to get rule 2");
    session
        .describe_data_policy_rule(policy_id.as_str(), rule_3_id.as_str())
        .expect("failed to get rule 3");

    session
        .update_data_policy_rule(
            policy_id.as_str(),
            rule_1_id.as_str(),
            NewDataPolicyRule {
                comment: None,
                clauses: vec![DataPolicyClause {
                    operator: Operator::AllOf,
                    tags: None,
                    capabilities: None,
                    facts: Some(vec![
                        FactExpression{
                            r#type: "fact_type".to_string(),
                            operator: antimatter_api::models::fact_expression::Operator::Exists,
                            arguments: vec![FactExpressionArgumentsInner{
                                operator: antimatter_api::models::fact_expression_arguments_inner::Operator::In,
                                values: Some(vec!["test".to_string()]),
                            }],
                            variables: None,
                        },
                    ]),
                    read_parameters: None,
                }],
                effect: DataPolicyRuleEffect::Allow,
                token_scope: Some(TokenScope::Capsule),
                token_format: Some(TokenFormat::Synthetic),
                priority: Some(5),
                assign_priority: None,
            },
        )
        .expect("failed to update data policy rule");

    let res = session
        .describe_data_policy_rule(policy_id.as_str(), rule_1_id.as_str())
        .expect("failed to get rule 1");
    assert_eq!(res.priority, 5, "unexpected data policy rule priority");

    let res = session
        .describe_data_policy(policy_id.as_str())
        .expect("failed to describe data policy");
    assert_eq!(
        res.rules.unwrap_or(vec![]).len(),
        3,
        "expected 3 data policy rules"
    );

    session
        .delete_data_policy_rule(policy_id.as_str(), rule_1_id.as_str())
        .expect("failed to delete data policy rule");

    let res = session
        .describe_data_policy(policy_id.as_str())
        .expect("failed to describe data policy");
    assert_eq!(
        res.rules.unwrap_or(vec![]).len(),
        2,
        "expected 2 data policy rules"
    );

    session
        .update_data_policy_rules(
            policy_id.as_str(),
            DataPolicyRuleChanges {
                new_rules: None,
                delete_rules: Some(vec![rule_2_id, rule_3_id]),
            },
        )
        .expect("failed to delete data policy rules");

    let res = session
        .describe_data_policy(policy_id.as_str())
        .expect("failed to describe data policy");
    assert_eq!(
        res.rules.unwrap_or(vec![]).len(),
        0,
        "expected 0 data policy rules"
    );

    session
        .delete_read_context("testreadctx1")
        .expect("failed to delete read context");

    let res = session
        .list_read_context()
        .expect("failed to list read contexts");
    assert_eq!(res.read_contexts.len(), 1, "expected 1 read context");
}

#[test]
fn test_api_policies() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    let res = session
        .list_policy_rules()
        .expect("failed to fetch policy rules");
    assert_eq!(res.rules.len(), 31, "expected 31 policies");

    let res = session
        .create_policy_rule(NewDomainPolicyRule {
            domain_identity: None,
            facts: Some(vec![FactPolicyRulesInner {
                operator: fact_policy_rules_inner::Operator::Exists,
                name: "fact1".to_string(),
                arguments: vec![FactPolicyRulesInnerArgumentsInner {
                    source: Any,
                    capability: Some("capability1".to_string()),
                    value: None,
                }],
            }]),
            path: "log".to_string(),
            operation: PolicyRuleOperation::View,
            result: PolicyRuleResult::Allow,
            priority: 123,
            disabled: true,
        })
        .expect("failed to create policy rules");
    assert_eq!(res.priority, 123, "unexpected priority");

    let resl = session
        .list_policy_rules()
        .expect("failed to fetch policy rules");
    assert_eq!(resl.rules.len(), 32, "expected 32 policies");

    session
        .update_policy_rule(
            res.id.as_str(),
            NewDomainPolicyRule {
                domain_identity: None,
                facts: Some(vec![FactPolicyRulesInner {
                    operator: fact_policy_rules_inner::Operator::Exists,
                    name: "fact1".to_string(),
                    arguments: vec![FactPolicyRulesInnerArgumentsInner {
                        source: Any,
                        capability: Some("capability1".to_string()),
                        value: None,
                    }],
                }]),
                path: "log".to_string(),
                operation: PolicyRuleOperation::View,
                result: PolicyRuleResult::Allow,
                priority: 123,
                disabled: true,
            },
        )
        .expect("failed to update policy rule");

    session
        .delete_policy_rule(res.id.as_str())
        .expect("failed to delete policy rule");

    let res = session
        .renumber_policy_rules()
        .expect("failed to renumber policy");
    assert_eq!(res.rules.len(), 31, "expected 31 policies");
}

#[test]
fn test_api_identity_providers() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    session
        .upsert_identity_providers(
            "testidp",
            DomainIdentityProviderDetails {
                google_o_auth: Some(Box::new(GoogleOAuthDomainIdentityProviderDetails {
                    client_id: None,
                    group_mappings: None,
                })),
                microsoft_o_auth: None,
                api_key: None,
            },
        )
        .expect("failed to create idp");

    let res = session
        .list_identity_providers()
        .expect("failed to list idp");
    assert_eq!(
        res.identity_providers.unwrap_or_default()[0].name,
        "apikey",
        "unexpected IDP name"
    );

    let res = session
        .get_identity_provider("testidp")
        .expect("failed to fetch idp 1");
    assert_eq!(
        res.supported_principals[0].to_string(),
        "Email",
        "unexpected IDP supported principal"
    );

    let _ = session
        .upsert_identity_providers(
            "testidp",
            DomainIdentityProviderDetails {
                google_o_auth: Some(Box::new(GoogleOAuthDomainIdentityProviderDetails {
                    client_id: Some("test.client.id".to_string()),
                    group_mappings: None,
                })),
                microsoft_o_auth: None,
                api_key: None,
            },
        )
        .expect("failed to upsert idp");

    let res = session
        .get_identity_provider("testidp")
        .expect("failed to fetch idp 2");

    let x = res.details.unwrap_or_default();
    assert_eq!(
        x.google_o_auth
            .unwrap_or_default()
            .client_id
            .unwrap_or_default(),
        "test.client.id",
        "unexpected IDP supported principal"
    );

    session
        .insert_identity_provider_principal(
            "testidp",
            DomainIdentityProviderPrincipalParams {
                capabilities: vec![],
                details: Box::new(DomainIdentityPrincipalDetails::Email(Box::new(
                    DomainIdentityEmailPrincipalParams {
                        r#type: Default::default(),
                        comment: Some("comment1".to_string()),
                        email: "test1@testing.com".to_string(),
                    },
                ))),
            },
        )
        .expect("failed to create idp principal 1");

    session
        .insert_identity_provider_principal(
            "testidp",
            DomainIdentityProviderPrincipalParams {
                capabilities: vec![],
                details: Box::new(DomainIdentityPrincipalDetails::Email(Box::new(
                    DomainIdentityEmailPrincipalParams {
                        r#type: Default::default(),
                        comment: Some("comment2".to_string()),
                        email: "test2@testing.com".to_string(),
                    },
                ))),
            },
        )
        .expect("failed to create idp principal 2");

    let res = session
        .get_identity_provider_principals("testidp")
        .expect("failed to fetch idp principals");
    assert_eq!(res.principals.len(), 2, "expected 2 principals");

    session
        .update_identity_provider_principal(
            "testidp",
            res.principals[0].principal_id.as_str(),
            UpdatePrincipalParams {
                capabilities: vec![Capability {
                    name: "admin".to_string(),
                    value: None,
                }],
                comment: None,
            },
        )
        .expect("failed to update idp principal");

    let resp = session
        .get_identity_provider_principal("testidp", res.principals[0].principal_id.as_str())
        .expect("failed to get idp principal");
    assert_eq!(resp.comment.len(), 8, "expected admin capability");

    session
        .delete_identity_provider_principal("testidp", res.principals[0].principal_id.as_str())
        .expect("failed to delete idp principal");

    let res = session
        .get_identity_provider_principals("testidp")
        .expect("failed to fetch idp principals");
    assert_eq!(res.principals.len(), 1, "expected 1 principals");

    session
        .delete_identity_provider("testidp")
        .expect("failed to delete idp");

    let res = session
        .list_identity_providers()
        .expect("failed to list idp");
    assert_eq!(
        res.identity_providers.unwrap_or_default().len(),
        3,
        "expected 3 idps"
    );
}

#[test]
fn test_idp_group_providers() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    let res = session
        .list_identity_group_providers()
        .expect("failed to get IDP group info");

    assert!(res.group_identity_providers.is_some());
    let providers = res.group_identity_providers.clone().unwrap();
    assert_eq!(providers.len(), 1);
    assert_eq!(
        providers[0].clone().name.unwrap(),
        "Google Domain Directory Group Membership Capability Mapping".to_string()
    );
}

#[test]
fn test_api_general() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    let res = session
        .get_private_info()
        .expect("failed to get private info");
    assert_eq!(
        res.default_display_name,
        session.get_domain_id(),
        "unexpected display name"
    );

    let res = session
        .get_public_info()
        .expect("failed to get public info");
    assert_eq!(res.id, session.get_domain_id(), "unexpected display name");

    let res = session.get_settings().expect("failed to fetch settings");
    assert_eq!(
        res.active_admin_contacts.unwrap_or_default()[0],
        "test@antimatter.io",
        "unexpected active email"
    );

    session.list_hooks().expect("failed to list hooks");

    let res = session.list_resources().expect("failed to list resources");
    assert_eq!(
        res.schema.len() > 10,
        true,
        "expected more than 10 resources"
    );

    let res = session
        .query_access_logs(None, None, None, None, None, None, None, None, None, None)
        .expect("failed to fetch access logs");
    assert_eq!(res.has_more, false, "expected no capsule access logs");

    let res = session
        .query_control_log(None, None, None, None, None, None, None)
        .expect("failed to fetch access logs");
    assert_eq!(
        res.results.len() > 0,
        true,
        "expected at least 1 control log"
    );
}

#[test]
fn test_api_facts() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    let res = session
        .list_fact_types()
        .expect("failed to list fact types 1");
    assert_eq!(res.fact_types.len(), 0, "expected 0 fact types");

    session
        .add_fact_type(
            "testtype1",
            NewFactTypeDefinition {
                description: "a test fact type 1".to_string(),
                arguments: vec![
                    NewFactTypeDefinitionArgumentsInner {
                        name: "name11".to_string(),
                        description: "desc11".to_string(),
                    },
                    NewFactTypeDefinitionArgumentsInner {
                        name: "name12".to_string(),
                        description: "desc12".to_string(),
                    },
                    NewFactTypeDefinitionArgumentsInner {
                        name: "name13".to_string(),
                        description: "desc13".to_string(),
                    },
                ],
            },
        )
        .expect("failed to add fact type 1");

    session
        .add_fact_type(
            "testtype2",
            NewFactTypeDefinition {
                description: "a test fact type 2".to_string(),
                arguments: vec![
                    NewFactTypeDefinitionArgumentsInner {
                        name: "name21".to_string(),
                        description: "desc21".to_string(),
                    },
                    NewFactTypeDefinitionArgumentsInner {
                        name: "name22".to_string(),
                        description: "desc22".to_string(),
                    },
                    NewFactTypeDefinitionArgumentsInner {
                        name: "name23".to_string(),
                        description: "desc23".to_string(),
                    },
                ],
            },
        )
        .expect("failed to add fact type 2");

    let res = session
        .list_fact_types()
        .expect("failed to list fact types 2");
    assert_eq!(res.fact_types.len(), 2, "expected 2 fact types");

    let res = session
        .get_fact_type("testtype1")
        .expect("failed to get fact type");
    assert_eq!(res.name, "testtype1", "unexpected fact type name");

    let res = session
        .list_facts("testtype2")
        .expect("failed to list fact types");
    assert_eq!(res.facts.len(), 0, "expected 0 facts of type");

    session
        .add_fact(
            "testtype1",
            NewFact {
                arguments: vec!["cake".to_string(), "mix".to_string(), "1".to_string()],
            },
        )
        .expect("failed to add fact 1");

    session
        .add_fact(
            "testtype2",
            NewFact {
                arguments: vec!["cake".to_string(), "mix".to_string(), "2".to_string()],
            },
        )
        .expect("failed to add fact 2");

    session
        .add_fact(
            "testtype1",
            NewFact {
                arguments: vec!["cake".to_string(), "mix".to_string(), "3".to_string()],
            },
        )
        .expect("failed to add fact 3");

    let res = session
        .list_facts("testtype2")
        .expect("failed to list fact types");
    assert_eq!(res.facts.len(), 1, "expected 0 facts of type");

    let res = session
        .get_fact("testtype2", res.facts[0].id.as_str())
        .expect("failed to get fact by id");
    assert_eq!(res.arguments[1], "mix", "unexpected fact argument");

    session
        .delete_fact(
            "testtype1",
            Some(vec!["cake".to_string(), "mix".to_string(), "1".to_string()]),
            None,
        )
        .expect("failed to delete fact by arguments");

    let res = session
        .list_facts("testtype1")
        .expect("failed to list fact types");
    assert_eq!(res.facts.len(), 1, "expected 0 facts of type");

    session
        .delete_fact("testtype1", None, Some(res.facts[0].id.as_str()))
        .expect("failed to delete fact by ID");

    let res = session
        .list_facts("testtype1")
        .expect("failed to list fact types");
    assert_eq!(res.facts.len(), 0, "expected 0 facts of type");

    session
        .delete_fact_type("testtype2", "testtype2")
        .expect("failed to delete fact type");

    let res = session
        .list_fact_types()
        .expect("failed to list fact types 2");
    assert_eq!(res.fact_types.len(), 1, "expected 1 fact types");

    session
        .add_fact(
            "testtype1",
            NewFact {
                arguments: vec!["cake".to_string(), "mix".to_string(), "1".to_string()],
            },
        )
        .expect("failed to add fact 1");

    session
        .add_fact(
            "testtype1",
            NewFact {
                arguments: vec!["cake".to_string(), "mix".to_string(), "2".to_string()],
            },
        )
        .expect("failed to add fact 2");

    session
        .add_fact(
            "testtype1",
            NewFact {
                arguments: vec!["cake".to_string(), "mix".to_string(), "3".to_string()],
            },
        )
        .expect("failed to add fact 3");

    session
        .add_fact(
            "testtype1",
            NewFact {
                arguments: vec!["cake".to_string(), "mix".to_string(), "5".to_string()],
            },
        )
        .expect("failed to add fact 5");

    session
        .add_fact(
            "testtype1",
            NewFact {
                arguments: vec!["cake".to_string(), "mix".to_string(), "4".to_string()],
            },
        )
        .expect("failed to add fact 4");

    let res = session
        .list_facts("testtype1")
        .expect("failed to list fact types 2");
    assert_eq!(res.facts.len(), 5, "expected 5 fact types");

    session
        .delete_fact_by_tuple(
            "testtype1",
            FactTuple {
                arguments: vec!["cake".to_string(), "mix".to_string(), "2".to_string()],
            },
        )
        .expect("failed delete a fact by tuple");

    let res = session
        .list_facts("testtype1")
        .expect("failed to list fact types 3");
    assert_eq!(res.facts.len(), 4, "expected 4 fact types");

    session
        .delete_fact(
            "testtype1",
            Some(vec!["cake".to_string(), "mix".to_string(), "4".to_string()]),
            None,
        )
        .expect("failed delete a fact by tuple 2");

    let res = session
        .list_facts("testtype1")
        .expect("failed to list fact types 4");
    assert_eq!(res.facts.len(), 3, "expected 3 fact types");

    session
        .delete_all_facts("testtype1")
        .expect("failed to delete all facts");

    let res = session
        .list_facts("testtype1")
        .expect("failed to list fact types 2");
    assert_eq!(res.facts.len(), 0, "expected 0 fact types");
}

// This endpoint was "not implemented" at time of writing.
#[test]
#[ignore]
fn test_api_encryption() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    session
        .flush_encryption_keys()
        .expect("failed to flush encryption keys");
}

#[test]
fn test_api_domains() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    session
        .new_peer_domain(
            Some(vec!["nick".to_string(), "name".to_string()]),
            None,
            "child1".to_string(),
            None,
            "child1display".to_string(),
            Some(true),
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
        )
        .expect("failed to create peer 1");

    session
        .new_peer_domain_link_all(
            Some(vec!["not".to_string(), "me".to_string()]),
            None,
            "child2".to_string(),
            None,
            "child2display".to_string(),
        )
        .expect("failed to create peer 2");

    let _ = session
        .get_peer(None, Some("child2"))
        .expect("failed to get peer 2");

    let res = session.list_peers().expect("failed to create get peers");

    assert_eq!(res.peers.len(), 2, "expected only 2 peers");

    let peer_id = &session.list_peers().expect("failure").peers[0].id;

    let peer_conf = session
        .get_peer_config(peer_id)
        .expect("failed to get peer 1 config");

    assert_eq!(
        peer_conf.display_name, "child1display",
        "unexpected display name pre update"
    );

    session
        .update_peer(
            None,
            Some("name"),
            None,
            "updateddisplayname".to_string(),
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            Some(42),
            None,
            None,
            None,
            None,
            None,
        )
        .expect("failed to update peer 1");

    let peer_conf = session
        .get_peer_config(peer_id)
        .expect("failed to get peer 1 config");

    assert_eq!(
        peer_conf.display_name, "updateddisplayname",
        "unexpected display name post update"
    );

    session
        .delete_peer(None, Some("me"), None)
        .expect("failed to delete peer 2");

    let res = session.list_peers().expect("failed to create get peers");

    assert_eq!(res.peers.len(), 1, "expected only 1 peer");

    session.get_top_tags().expect("failed to get top tags");

    let err = session
        .get_peer(None, Some("unknown"))
        .err()
        .expect("expected an error");
    match err {
        SessionError::Status404(e) => {}
        _ => panic!("unexpected error type (expected 404): {:?}", err),
    }
}

#[test]
fn test_api_capsules() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    let input_data = vec![vec!["row0, col0: example data element".as_bytes().to_vec()]];

    let cell_readers = convert_to_readers(input_data.clone(), vec![vec![vec![]]])
        .expect("failed to generate data");

    let columns = vec![Column {
        name: "col0".to_string(),
        tags: vec![],
        skip_classification: false,
    }];

    let cfg = EncapsulateConfig {
        write_context_name: "default".to_string(),
        extra: "some extra data".to_string(),
        subdomain: None,
        subdomain_from: None,
        create_subdomains: None,
        async_seal: false,
    };

    // encapsulate the data
    let (mut reader, _meta) = session
        .encapsulate(columns, cell_readers, vec![], cfg)
        .expect("failed to encapsulate");
    let mut capsule_data: Vec<u8> = Vec::new();
    reader
        .read_to_end(&mut capsule_data)
        .expect("failed to read capsule data");
    // reader seems to be linked to the session so if we don't drop it, we get
    // Err: cannot borrow `session` as mutable more than once at a time
    drop(reader);

    // open the sealed capsule
    let writer = Cursor::new(capsule_data);
    let mut iterator = session
        .open("default", HashMap::new(), HashMap::new(), writer)
        .expect("failed to open capsule");
    let (_tags, output_data) = iterator.read_all(&[]).expect("failed to read all");

    assert_eq!(input_data, output_data);

    // ---

    let res_list = session
        .list_capsules(None, None, None, None, None, None, None)
        .expect("failed to fetch capsule list");

    assert_eq!(res_list.results.len(), 1, "expected only 1 capsule");

    let res_info = session
        .get_capsule_info(res_list.results[0].id.as_str())
        .expect("failed to fetch capsule list");

    assert_eq!(
        res_info.capsule_tags.len(),
        0,
        "expected 0 capsule tags pre delete"
    );

    let cap_tags = vec![models::Tag {
        name: "name1".to_string(),
        value: "value1".to_string(),
        r#type: TagTypeField::String,
        source: "source1".to_string(),
        hook_version: Some("1.0.0".to_string()),
    }];
    session
        .upsert_capsule_tags(res_list.results[0].id.as_str(), cap_tags)
        .expect("failed to upsert capsule tags");

    let res_info = session
        .get_capsule_info(res_list.results[0].id.as_str())
        .expect("failed to fetch capsule list");

    assert_eq!(res_info.capsule_tags.len(), 1, "expected 1 capsule tags");

    session
        .delete_capsule_tags(
            res_list.results[0].id.as_str(),
            DeleteTags {
                names: Some(vec!["name1".to_string()]),
            },
        )
        .expect("failed to delete capsule tags");

    let res_info = session
        .get_capsule_info(res_list.results[0].id.as_str())
        .expect("failed to fetch capsule list");

    assert_eq!(
        res_info.capsule_tags.len(),
        0,
        "expected 0 capsule tags post delete"
    );
}

#[test]
fn test_domain_capabilities() {
    let res = domains::create_domain("test@antimatter.io").expect("failed to create session");
    let mut session = res.0;

    session
        .put_capability(
            "cap1",
            "summary1".to_string(),
            Some("desc1".to_string()),
            true,
            Some(true),
        )
        .expect("failed to create capability 1");

    session
        .put_capability(
            "cap2",
            "summary2".to_string(),
            Some("desc2".to_string()),
            false,
            None,
        )
        .expect("failed to create capability 2");

    let res = session
        .get_capability("cap1")
        .expect("failed to fetch capability 1");

    assert_eq!(res.summary, "summary1", "unexpected description");

    let list_res = session
        .list_capabilities()
        .expect("failed to fetch capability list pre delete");

    assert_eq!(
        list_res.capabilities.len(),
        5,
        "expected only 2 capabilities"
    );

    session
        .delete_capability("cap1")
        .expect("failed to delete capability");

    let list_res = session
        .list_capabilities()
        .expect("failed to fetch capability list post delete");

    assert_eq!(
        list_res.capabilities.len(),
        4,
        "expected only 1 capabilities"
    );
}

fn convert_to_readers(
    elements: Vec<Vec<Vec<u8>>>,
    tags: Vec<Vec<Vec<SpanTag>>>,
) -> Result<Vec<RowReader>, String> {
    if elements.is_empty() {
        return Ok(Vec::new());
    }

    let col_count = elements[0].len();
    // Validate column consistency.
    if elements.iter().any(|row| row.len() != col_count) {
        return Err("column length inconsistency".to_string());
    }

    let rows = elements
        .clone()
        .into_iter()
        .zip(tags.into_iter())
        .map(|(row, tags)| {
            let mapped = row
                .clone()
                .into_iter()
                .zip(tags.clone().into_iter())
                .map(|(item_a, item_b)| to_data_element(item_a, item_b))
                .collect::<Result<Vec<CellReader>, String>>()
                .unwrap();
            RowReader {
                tags: vec![],
                cells: mapped,
            }
        })
        .collect::<Vec<RowReader>>();
    Ok(rows)
}
fn to_data_element(element: Vec<u8>, tags: Vec<SpanTag>) -> Result<CellReader, String> {
    CellReader::new(tags, std::io::Cursor::new(element.clone()))
        .map_err(|e| format!("failed to create reader for element: {}", e))
}